pybaum.typecheck#

Module Contents#

Functions#

get_type(obj)

namdetuple aware type check.

pybaum.typecheck.get_type(obj)[source]#

namdetuple aware type check.

As in JAX we treat collections.namedtuple and typing.NamedTuple both as namedtuple but the exact type is preserved in the unflatten function.

namedtuples are discovered by being instances of tuple and having a _fields attribute as suggested by Raymond Hettinger here.

Moreover we check for the presence of a _replace method because we need when unflattening pytrees.

This can produce false positives but in most cases would still result in desired behavior.

Parameters

obj – The object to be checked

Returns

bool