Home
last modified time | relevance | path

Searched refs:needs_pytree (Results 1 – 6 of 6) sorted by relevance

/external/pytorch/torch/distributed/tensor/
D_op_schema.py210 needs_pytree: bool = False variable in RuntimeSchemaInfo
247 if self.schema_info is not None and self.schema_info.needs_pytree
259 if self.schema_info is not None and self.schema_info.needs_pytree
411 and origin_schema.schema_info.needs_pytree
D_dispatch.py339 if runtime_schema_info is not None and runtime_schema_info.needs_pytree:
/external/pytorch/torch/distributed/tensor/experimental/
D_register_sharding.py123 static_argnum, static_kwargkey or None, needs_pytree=True
/external/pytorch/torch/distributed/tensor/_ops/
D_tensor_ops.py503 @register_op_strategy(aten.stack.default, RuntimeSchemaInfo(1, needs_pytree=True))
538 @register_op_strategy(aten.cat.default, RuntimeSchemaInfo(1, needs_pytree=True))
605 @register_prop_rule(aten.index.Tensor, schema_info=RuntimeSchemaInfo(needs_pytree=True))
D_pointwise_ops.py665 register_op_strategy(op, schema_info=RuntimeSchemaInfo(needs_pytree=True))(
670 register_op_strategy(op, schema_info=RuntimeSchemaInfo(needs_pytree=True))(
686 register_op_strategy(op, schema_info=RuntimeSchemaInfo(needs_pytree=True))(
D_math_ops.py398 [aten._foreach_norm.Scalar], schema_info=RuntimeSchemaInfo(1, needs_pytree=True)