# TODO(mdan): Index more efficiently. Could do a name check instead.
    if any(v is value for v in AG_MODULE.__dict__.values()):
      return {TFRTypes.AG_BUILTIN_FUNC}
    if getattr(value, '__name__', None) == 'tensorflow.raw_ops':
      return {types.ModuleType}
    if hasattr(value, '__module__'):
      if isinstance(value, dtypes.DType):
        return {TFRTypes.ATTR}

  def model_fn(args):
    # do something with dynamic tensor

  @function.Defun(capture_resource_var_by_value=False)
  def tpu_subgraph():
      return tf.tpu.rewrite(model_fn, args)

  return tf.raw_ops.TPUPartitionedCall(
      args=tpu_subgraph.captured_inputs,
      Tout=[o.type for o in tpu_subgraph.definition.signature.output_arg],
      f=tpu_subgraph,
      device_ordinal=[0])