[pytree] support PyStructSequence types for Python pytree #113258

zou3519 · 2023-11-30T16:21:34Z

@XuehaiPan: what's the motivation for having all PyStructSequence be pytrees? All namedtuples being pytrees is already unexpected behavior for some people.

what's the motivation for having all PyStructSequence be pytrees?

@zou3519 The idea is there might be a use case that someone writes a custom operator and wants it supported by pytree by default. Currently, if a third-party package has a custom return type, it needs to call torch.return_types.pytree_register_structseq manually.

There is also a codegen utility to generate C++ code for PyStructSequence type and assign it to torch._C._return_types.

pytorch/tools/autograd/gen_python_functions.py

Lines 751 to 768 in 2a3d8e5

if typename is None:

typename = f'{name}NamedTuple{"" if not definitions else len(definitions)}'

typenames[tn_key] = typename

definitions.append(

f"""\

PyTypeObject* get_{name}_namedtuple() {{

static PyStructSequence_Field NamedTuple_fields[] = {{ {fields}, {{nullptr}} }};

static PyTypeObject {typename};

static bool is_initialized = false;

static PyStructSequence_Desc desc = {{ "torch.return_types.{name}", nullptr, NamedTuple_fields, {len(fieldnames)} }};

if (!is_initialized) {{

PyStructSequence_InitType(&{typename}, &desc);

{typename}.tp_repr = (reprfunc)torch::utils::returned_structseq_repr;

is_initialized = true;

}}

return &{typename};

}}

"""

See also:

torchgen/api/python.py

tools/autograd/gen_python_functions.py

tools/autograd/templates/python_return_types.h

tools/autograd/templates/python_return_types.h

All namedtuples being pytrees is already unexpected behavior for some people.

I agree that some people want their namedtuple class to be an opaque object rather than a container. Using dataclasses.dataclass might be an alternative for them. Before this PR, all namedtuples are already pytrees and all dataclasses are leaf nodes by default.

See also discussion for making all dataclasses be pytrees by default:

[export] allow register dataclass as pytree node #106160

Add dataclass support to pytree #93214 (comment)

Automatically treat dataclasses as pytrees jax-ml/jax#2371 (comment)

The idea is there might be a use case that someone writes a custom operator and wants it supported by pytree by default. Currently, if a third-party package has a custom return type, it needs to call torch.return_types.pytree_register_structseq manually.

What does JAX's pytree do on structseq types?

For custom ops that return structseq: I don't think there is a way to easily do this, so I wouldn't worry about that use case.

What does JAX's pytree do on structseq types?

By default, structseqs are leaf nodes and all namedtuples are pytree nodes in JAX's pytree. Users need to manually register structseq types are pytree nodes like what we do for pytree_register_structseq.

The JAX community uses namedtuples instead of structseqs. They never use structseqs. If you search PyStructSequence in https://github.com/google/jax and https://github.com/tensorflow/tensorflow, you will get an empty result.

https://github.com/search?q=repo%3Agoogle%2Fjax+PyStructSequence&type=code
https://github.com/search?q=repo%3Atensorflow%2Ftensorflow+PyStructSequence&type=code

Here are lists of structseq types:

CPython internals

In [1]: import sys In [2]: import importlib In [3]: for mod in sys.stdlib_module_names: ...: try: ...: importlib.import_module(mod) ...: except ImportError: ...: pass ...: In [4]: import optree In [5]: list(filter(optree.is_structseq_class, tuple.__subclasses__())) Out[5]: [ <class 'sys.int_info'>, <class 'sys.float_info'>, <class 'UnraisableHookArgs'>, <class 'sys.hash_info'>, <class 'sys.version_info'>, <class 'sys.flags'>, <class 'sys.thread_info'>, <class 'asyncgen_hooks'>, <class '_thread._ExceptHookArgs'>, <class 'os.stat_result'>, <class 'os.statvfs_result'>, <class 'os.terminal_size'>, <class 'posix.times_result'>, <class 'posix.uname_result'>, <class 'time.struct_time'>, <class 'resource.struct_rusage'>, <class '_lsprof.profiler_entry'>, <class '_lsprof.profiler_subentry'>, <class 'curses.ncurses_version'>, <class 'pwd.struct_passwd'>, <class 'grp.struct_group'> ]

With `torch.return_types`

In [6]: import torch In [7]: list(filter(optree.is_structseq_class, tuple.__subclasses__())) Out[7]: [ <class 'sys.int_info'>, <class 'sys.float_info'>, <class 'UnraisableHookArgs'>, <class 'sys.hash_info'>, <class 'sys.version_info'>, <class 'sys.flags'>, <class 'sys.thread_info'>, <class 'asyncgen_hooks'>, <class '_thread._ExceptHookArgs'>, <class 'os.stat_result'>, <class 'os.statvfs_result'>, <class 'os.terminal_size'>, <class 'posix.times_result'>, <class 'posix.uname_result'>, <class 'time.struct_time'>, <class 'resource.struct_rusage'>, <class '_lsprof.profiler_entry'>, <class '_lsprof.profiler_subentry'>, <class 'curses.ncurses_version'>, <class 'pwd.struct_passwd'>, <class 'grp.struct_group'>, <class 'numpy.core.multiarray.typeinfo'>, <class 'numpy.core.multiarray.typeinforanged'>, <class 'torch.return_types._fake_quantize_per_tensor_affine_cachemask_tensor_qparams'>, <class 'torch.return_types._fused_moving_avg_obs_fq_helper'>, <class 'torch.return_types._linalg_det'>, <class 'torch.return_types._linalg_det_out'>, <class 'torch.return_types._linalg_eigh'>, <class 'torch.return_types._linalg_eigh_out'>, <class 'torch.return_types._linalg_slogdet'>, <class 'torch.return_types._linalg_slogdet_out'>, <class 'torch.return_types._linalg_solve_ex'>, <class 'torch.return_types._linalg_solve_ex_out'>, <class 'torch.return_types._linalg_svd'>, <class 'torch.return_types._linalg_svd_out'>, <class 'torch.return_types._lu_with_info'>, <class 'torch.return_types._scaled_dot_product_efficient_attention'>, <class 'torch.return_types._scaled_dot_product_flash_attention'>, <class 'torch.return_types._unpack_dual'>, <class 'torch.return_types.aminmax'>, <class 'torch.return_types.aminmax_out'>, <class 'torch.return_types.cummax'>, <class 'torch.return_types.cummax_out'>, <class 'torch.return_types.cummin'>, <class 'torch.return_types.cummin_out'>, <class 'torch.return_types.frexp'>, <class 'torch.return_types.frexp_out'>, <class 'torch.return_types.geqrf_out'>, <class 'torch.return_types.geqrf'>, <class 'torch.return_types.histogram_out'>, <class 'torch.return_types.histogram'>, <class 'torch.return_types.histogramdd'>, <class 'torch.return_types.kthvalue'>, <class 'torch.return_types.kthvalue_out'>, <class 'torch.return_types.linalg_cholesky_ex'>, <class 'torch.return_types.linalg_cholesky_ex_out'>, <class 'torch.return_types.linalg_eig'>, <class 'torch.return_types.linalg_eig_out'>, <class 'torch.return_types.linalg_eigh'>, <class 'torch.return_types.linalg_eigh_out'>, <class 'torch.return_types.linalg_inv_ex'>, <class 'torch.return_types.linalg_inv_ex_out'>, <class 'torch.return_types.linalg_ldl_factor'>, <class 'torch.return_types.linalg_ldl_factor_out'>, <class 'torch.return_types.linalg_ldl_factor_ex'>, <class 'torch.return_types.linalg_ldl_factor_ex_out'>, <class 'torch.return_types.linalg_lstsq'>, <class 'torch.return_types.linalg_lstsq_out'>, <class 'torch.return_types.linalg_lu'>, <class 'torch.return_types.linalg_lu_out'>, <class 'torch.return_types.linalg_lu_factor'>, <class 'torch.return_types.linalg_lu_factor_out'>, <class 'torch.return_types.linalg_lu_factor_ex'>, <class 'torch.return_types.linalg_lu_factor_ex_out'>, <class 'torch.return_types.linalg_qr'>, <class 'torch.return_types.linalg_qr_out'>, <class 'torch.return_types.linalg_slogdet'>, <class 'torch.return_types.linalg_slogdet_out'>, <class 'torch.return_types.linalg_solve_ex'>, <class 'torch.return_types.linalg_solve_ex_out'>, <class 'torch.return_types.linalg_svd'>, <class 'torch.return_types.linalg_svd_out'>, <class 'torch.return_types.lu_unpack'>, <class 'torch.return_types.lu_unpack_out'>, <class 'torch.return_types.max'>, <class 'torch.return_types.max_out'>, <class 'torch.return_types.median'>, <class 'torch.return_types.median_out'>, <class 'torch.return_types.min'>, <class 'torch.return_types.min_out'>, <class 'torch.return_types.mode'>, <class 'torch.return_types.mode_out'>, <class 'torch.return_types.nanmedian'>, <class 'torch.return_types.nanmedian_out'>, <class 'torch.return_types.qr_out'>, <class 'torch.return_types.qr'>, <class 'torch.return_types.slogdet'>, <class 'torch.return_types.slogdet_out'>, <class 'torch.return_types.sort_out'>, <class 'torch.return_types.sort'>, <class 'torch.return_types.svd_out'>, <class 'torch.return_types.svd'>, <class 'torch.return_types.topk_out'>, <class 'torch.return_types.topk'>, <class 'torch.return_types.triangular_solve_out'>, <class 'torch.return_types.triangular_solve'> ]

For the use cases for pytree in PyTorch, the leaf values are mostly torch.Tensor or None. The Python internal PyStructSequeunce types should never be used as a container in the PyTorch use case.

IMO we should be consistent, so we should not register structseq as pytree nodes by default. However, users can manually add them (and we "manually" add them for all of the torch.return_types).

Does optree support structseq types as pytree nodes by default? If so, we probably want some way to toggle this behavior to make it consistent.

IMO we should be consistent, so we should not register structseq as pytree nodes by default. However, users can manually add them (and we "manually" add them for all of the torch.return_types).

@zou3519 Do you mean to be consistent with JAX pytree or the old behavior of pytree in PyTorch?

Does optree support structseq types as pytree nodes by default?

Yes. The motivations are:

PyStructSequence is a concept in Python C API but shares some similar semantics of namedtuple.

Python Documentation: Struct Sequence Objects

Struct sequence objects are the C equivalent of namedtuple() objects, i.e. a sequence whose items can also be accessed through attributes. To create a struct sequence, you first have to create a specific struct sequence type.

Also, the CPython upstream may want to add members to PyStructSequence to make it consistent with namedtuple in Python side.

Enhance Object/structseq.c to match namedtuple and tuple api python/cpython#46145

GH-46145: make PyStructSequence compatible with namedtuple python/cpython#108648

Since we treat all namedtuple types are pytree nodes, it is reasonable to have all PyStructSequence types be pytree nodes if the upstream PR is landed.

Here is the implementation for is_namedtuple_class in the previous PR in the ghstack:

def is_namedtuple_class(cls): return ( isinstance(cls, type) and issubclass(cls, tuple) and isinstance(getattr(cls, "_fields", None), tuple) and all(type(field) is str for field in cls._fields) and callable(getattr(cls, "_make", None)) # <= structseq fails here even if python/cpython#108648 landed and callable(getattr(cls, "_asdict", None)) )

We will get is_namedtuple(structseq) to be False even if GH-46145: make PyStructSequence compatible with namedtuple python/cpython#108648 landed.

But for JAX pytree, it only checks the presence of _fields:

def is_namedtuple(obj): return isinstance(obj, tuple) and hasattr(obj, '_fields')

All PyStructSequence instances will be namedtuple if GH-46145: make PyStructSequence compatible with namedtuple python/cpython#108648 landed. I.e., they will be pytree nodes automatically in JAX pytree.

PyStructSequence types are subclasses of tuple, and they are container types. This fits the definition of pytree node. That's why we manually register all types in torch.return_types.* as pytree node. Here is an example usage that requires this: tree_map(lambda x: x.cuda(), tree).

The JAX community does not use PyStructSequence.

The number of PyStructSequence is limited (see [pytree] support PyStructSequence types for Python pytree #113258 (comment)). Because it is related to C code and there is no easy way to create a new type from the Python side.

If so, we probably want some way to toggle this behavior to make it consistent.

I can add an API to have namespace-specific configurations in optree. But before that, I want to hear your thoughts about the explanations above.

@zou3519 Do you mean to be consistent with JAX pytree or the old behavior of pytree in PyTorch?

We should be consistent with JAX pytree. Also, the old behavior of pytree in PyTorch is the same as JAX's behavior: structseq is not treated as a pytree node type, unless the user explicitly specifies it. We've explicitly specified that all of the torch.return_types are pytree nodes.

PyStructSequence is a concept in Python C API but shares some similar semantics of namedtuple.

The only user of structseq we've seen out in the wild is torch.return_types. So it doesn't matter too much what we do with it, but we should be consistent because inconsistencies end up leading to UX issues.

All PyStructSequence instances will be namedtuple if landed

I would wait until those PRs actually land before changing the behavior, to be consistent

PyStructSequence types are subclasses of tuple, and they are container types. This fits the definition of pytree node. That's why we manually register all types in torch.return_types.* as pytree node. Here is an example usage that requires this: tree_map(lambda x: x.cuda(), tree).

I don't think subclasses of a type get the pytree treatment; we require exact match, right?
I agree PyStructSequence is a container type, but not all container types are pytree nodes by default: see dataclass.

We manually register all types in torch.return_types.* as pytree nodes because we own them and as the owners, we've decided they should have pytree semantics.

All PyStructSequence instances will be namedtuple if landed
I would wait until those PRs actually land before changing the behavior, to be consistent

Should we ask on the PR when they think they'll merge it? lol

Should we ask on the PR when they think they'll merge it? lol

Based on the label of the issue, the initial plan is to have this in Python 3.9 but delayed to 3.13. I think they will merge relevant changes eventually.

Enhance Object/structseq.c to match namedtuple and tuple api python/cpython#46145

Make structseq more like collections.namedtuple python/cpython#46561

GH-46145: make PyStructSequence compatible with namedtuple python/cpython#108648

Maybe we should just leave this PR open and avoid changing optree for now. As we discussed above, people don't really use the structseq API.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[pytree] support PyStructSequence types for Python pytree #113258

Diff view

Diff view

Uh oh!

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

[pytree] support PyStructSequence types for Python pytree #113258

Are you sure you want to change the base?

[pytree] support PyStructSequence types for Python pytree #113258

Uh oh!

Uh oh!

Diff view

Diff view

Uh oh!

There are no files selected for viewing

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!