def signature(obj):
  """Returns a `ShapeDtype` signature for the given `obj`.

  A signature is either a `ShapeDtype` instance or a tuple of `ShapeDtype`
  instances. Note that this function is permissive with respect to its inputs
  (accepts lists or tuples or dicts, and underlying objects can be any type
  as long as they have shape and dtype attributes) and returns the corresponding
  nested structure of `ShapeDtype`.

  Args:
    obj: An object that has `shape` and `dtype` attributes, or a list/tuple/dict
        of such objects.

  Returns:
    A corresponding nested structure of `ShapeDtype` instances.
  """
  if isinstance(obj, (list, tuple)):
    output = tuple(signature(x) for x in obj)
    return output if isinstance(obj, tuple) else list(output)
  elif isinstance(obj, dict):
    return {k: signature(v) for (k, v) in obj.items()}
  else:
    return ShapeDtype(obj.shape, obj.dtype)
Returns a `ShapeDtype` signature for the given `obj`.

  A signature is either a `ShapeDtype` instance or a tuple of `ShapeDtype`
  instances. Note that this function is permissive with respect to its inputs
  (accepts lists or tuples or dicts, and underlying objects can be any type
  as long as they have shape and dtype attributes) and returns the corresponding
  nested structure of `ShapeDtype`.

  Args:
    obj: An object that has `shape` and `dtype` attributes, or a list/tuple/dict
        of such objects.

  Returns:
    A corresponding nested structure of `ShapeDtype` instances.
  
trax/shapes.py