Verify model supports segment-ids

[yatlas20098]

• I have added all the necessary unit tests for my change.
• I have verified that my change does not break existing code and all unit tests pass.
• I have added all appropriate doc-strings/documentation.
• My PR is based on the latest changes of the main branch (if unsure, rebase the code).
• I have signed the Contributor License Agreement.
• I have followed Contribution Guidelines.

[tianshub]

worth to move it out as a helper function and cache it? so it's not part of the core logic.

e.g.

@functools.cache
def call_contains(model, target_arg):
  ...

[yatlas20098]

I've implemented caching but key'd it based on the model type instead of the model instance. Since nnx.merge constructs the model instance during each computation step, caching the model instance directly would result in zero hits.