-
Notifications
You must be signed in to change notification settings - Fork 3
Expand file tree
/
Copy patharray_dataclass.py
More file actions
1253 lines (1038 loc) · 42 KB
/
Copy patharray_dataclass.py
File metadata and controls
1253 lines (1038 loc) · 42 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# Copyright 2026 The dataclass_array Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Dataclass array."""
from __future__ import annotations
import dataclasses
import functools
import sys
import typing
from typing import Any, Callable, ClassVar, Generic, Iterator, Optional, Set, Tuple, Type, TypeVar, Union
from dataclass_array import field_utils
from dataclass_array import shape_parsing
from dataclass_array import type_parsing
from dataclass_array.typing import Array
from dataclass_array.typing import Axes, DTypeArg, DcOrArray, DcOrArrayT, DynamicShape, Shape # pylint: disable=g-multiple-import,g-importing-member
from dataclass_array.utils import np_utils
from dataclass_array.utils import py_utils
import einops
from etils import array_types
from etils import edc
from etils import enp
from etils import epy
import numpy as np
import typing_extensions
from typing_extensions import Annotated, Literal, TypeAlias # pylint: disable=g-multiple-import
lazy = enp.lazy
# TODO(pytype): Should use `dca.typing.DcT` but bound does not work across
# modules.
_DcT = TypeVar('_DcT', bound='DataclassArray')
# Any valid numpy indices slice ([x], [x:y], [:,...], ...)
_IndiceItem = Union[type(Ellipsis), None, int, slice, Any]
_Indices = Tuple[_IndiceItem] # Normalized slicing
_IndicesArg = Union[_IndiceItem, _Indices]
_METADATA_KEY = 'dca_field'
_DUMMY_ARRAY_FIELD = '_dca_dummy_array'
@dataclasses.dataclass(frozen=True)
class DataclassParams:
"""Params controlling the DataclassArray behavior.
Set by `@dca.dataclass_array`. Saved in `cls.__dca_params__`.
Attributes:
broadcast: If `True`, enable input broadcasting
cast_dtype: If `True`, auto-cast inputs `dtype`
cast_list: If `True`, auto-cast lists to `xnp.ndarray`
"""
# If modifying this, make sure to modify `@dataclass_array` too!
broadcast: bool = False
cast_dtype: bool = False
cast_list: bool = True
def dataclass_array(
*,
# If modifying this, make sure to modify `DataclassParams` too!
broadcast: bool = False,
cast_dtype: bool = False,
cast_list: bool = True,
) -> Callable[[type[_DcT]], type[_DcT]]:
"""Optional decorator to customize `dca.DataclassArray` params.
Usage:
```python
@dca.dataclass_array()
class MyDataclass(dca.DataclassArray):
...
```
This decorator has to be added in addition of inheriting from
`dca.DataclassArray`.
Args:
broadcast: If `True`, enable input broadcasting
cast_dtype: If `True`, auto-cast inputs `dtype`
cast_list: If `True`, auto-cast lists to `xnp.ndarray`
Returns:
decorator: The decorator which will apply the options to the dataclass
"""
def decorator(cls):
if not issubclass(cls, DataclassArray):
raise TypeError(
'`@dca.dataclass_array` can only be applied on `dca.DataclassArray`. '
f'Got: {cls}'
)
cls.__dca_params__ = DataclassParams(
broadcast=broadcast,
cast_dtype=cast_dtype,
cast_list=cast_list,
)
return cls
return decorator
def array_field(
shape: Shape,
dtype: DTypeArg = float,
**field_kwargs,
) -> dataclasses.Field[DcOrArray]:
"""Dataclass array field.
See `dca.DataclassArray` for example.
Args:
shape: Inner shape of the field
dtype: Type of the field
**field_kwargs: Args forwarded to `dataclasses.field`
Returns:
The dataclass field.
"""
# TODO(epot): Validate shape, dtype
dca_field = _ArrayFieldMetadata(
inner_shape_non_static=shape,
dtype=dtype, # pyrefly: ignore[bad-argument-type]
)
return dataclasses.field(**field_kwargs, metadata={_METADATA_KEY: dca_field})
class MetaDataclassArray(type):
"""DataclassArray metaclass."""
# TODO(b/204422756): We cannot use `__class_getitem__` due to b/204422756
def __getitem__(cls, spec):
# Not clear how this would interact if cls is also a `Generic`
return Annotated[cls, field_utils.ShapeAnnotation(spec)] # pyrefly: ignore[not-a-type]
# TODO(epot): Restore once pytype support this
# @typing_extensions.dataclass_transform(
# kw_only_default=True,
# # TODO(b/272524683):Restore field specifier
# # field_specifiers=(
# # dataclasses.Field,
# # dataclasses.field,
# # array_field,
# # ),
# )
class DataclassArray(metaclass=MetaDataclassArray):
"""Dataclass which behaves like an array.
Usage:
```python
class Square(DataclassArray):
pos: f32['*shape 2']
scale: f32['*shape']
name: str
# Create 3 squares batched
p = Square(
pos=[[x0, y0], [x1, y1], [x2, y2]],
scale=[scale0, scale1, scale2],
name='my_square',
)
p.shape == (3,)
p.pos.shape == (3, 2)
p[0] == Square(pos=[x0, y0], scale=scale0)
p = p.reshape((3, 1)) # Reshape the inner-shape
p.shape == (3, 1)
p.pos.shape == (3, 1, 2)
p.name == 'my_square'
```
`DataclassArray` has 2 types of fields:
* Array fields: Fields batched like numpy arrays, with reshape, slicing,...
(`pos` and `scale` in the above example).
* Static fields: Other non-numpy field. Are not modified by reshaping,... (
`name` in the above example).
Static fields are also ignored in `jax.tree.map`.
`DataclassArray` detect array fields if either:
* The typing annotation is a `etils.array_types` annotation (in which
case shape/dtype are automatically infered from the typing annotation)
Example: `x: f32[..., 3]`
* The typing annotation is another `dca.DataclassArray` (in which case
`my_dataclass.field.shape == my_dataclass.shape`)
Example: `x: MyDataclass`
* The field is explicitly defined in `dca.array_field`, in which case
the typing annotation is ignored.
Example: `x: Any = dca.field(shape=(), dtype=np.int64)`
Field which do not satisfy any of the above conditions are static (including
field annotated with `field: np.ndarray` or similar).
"""
# Child class inherit the default params by default, but can also
# overwrite them.
__dca_params__: ClassVar[DataclassParams] = DataclassParams()
# TODO(epot): Could be removed with py3.10 and using `kw_only=True`
# Fields defined here will be forwarded with `.replace`
# TODO(py39): Replace Set -> set
__dca_non_init_fields__: ClassVar[Set[str]] = set()
_shape: Shape
_xnp: enp.NpModule
def __init_subclass__(
cls,
frozen=True,
**kwargs,
):
super().__init_subclass__(**kwargs)
if not frozen:
raise ValueError(f'{cls} cannot be `frozen=False`.')
# Apply dataclass (in-place)
if not typing.TYPE_CHECKING:
# TODO(b/227290126): Create pytype issues
dataclasses.dataclass(frozen=True)(cls)
# TODO(epot): Could have smart __repr__ which display types if array have
# too many values (maybe directly in `edc.field(repr=...)`).
edc.dataclass(kw_only=True, repr=True, auto_cast=False)(cls)
cls._dca_jax_tree_registered = False
cls._dca_torch_tree_registered = False
# Typing annotations have to be lazily evaluated (to support
# `from __future__ import annotations` and forward reference)
# To avoid costly `typing.get_type_hints` which perform `eval` and `str`
# convertions, we cache the type annotations here.
cls._dca_fields_metadata: Optional[dict[str, _ArrayFieldMetadata]] = None
# Normalize the `cls.__dca_non_init_fields__`
# TODO(epot): Support inheritance if the parents also define
# `__dca_non_init_fields__` (fields should be merged from `.mro()`)
cls.__dca_non_init_fields__ = set(cls.__dca_non_init_fields__)
if typing.TYPE_CHECKING:
# TODO(b/242839979): pytype do not support PEP 681 -- Data Class Transforms
def __init__(self, **kwargs):
pass
def __post_init__(self) -> None:
"""Validate and normalize inputs."""
cls = type(self)
# First time, we perform additional check & updates
if cls._dca_fields_metadata is None: # pylint: disable=protected-access
_init_cls(self)
# Register the tree_map here instead of `__init_subclass__` as `jax` may
# not have been imported yet during import.
if enp.lazy.has_jax and not cls._dca_jax_tree_registered: # pylint: disable=protected-access
enp.lazy.jax.tree_util.register_pytree_with_keys_class(cls)
cls._dca_jax_tree_registered = True # pylint: disable=protected-access
if enp.lazy.has_torch and not cls._dca_torch_tree_registered: # pylint: disable=protected-access
# Note: Torch is updating it's tree API to make it public and use `optree`
# as backend: https://github.com/pytorch/pytorch/issues/65761
enp.lazy.torch.utils._pytree._register_pytree_node( # pylint: disable=protected-access
cls,
flatten_fn=lambda a: a._tree_flatten(), # pylint: disable=protected-access
unflatten_fn=lambda vals, ctx: cls.tree_unflatten(ctx, vals),
)
cls._dca_torch_tree_registered = True # pylint: disable=protected-access
# Validate and normalize array fields
# * Maybe cast (list, np) -> xnp
# * Maybe cast dtype
# * Maybe broadcast shapes
# Because this is only done inside `__init__`, it is ok to mutate self.
# Cast and validate the array xnp are consistent
xnp = self._cast_xnp_dtype_inplace()
# Validate the batch shape is consistent
# However, we need to be careful that `_ArrayField` never uses
# `@functools.cached_property`
shape = self._broadcast_shape_inplace()
# TODO(epot): When to validate (`field.validate()`)
if xnp is None: # No values
# Inside `jax.tree_utils`, tree-def can be created with `None` values.
# Inside `jax.vmap`, tree can be created with `object()` sentinel values.
assert shape is None
xnp = None
# Cache results
# Should the state be stored in a separate object to avoid collisions ?
assert shape is None or isinstance(shape, tuple), shape
self._setattr('_shape', shape)
self._setattr('_xnp', xnp)
# ====== Array functions ======
@property
def shape(self) -> Shape:
"""Returns the batch shape common to all fields."""
return self._shape
@property
def size(self) -> int:
"""Returns the number of elements."""
return np_utils.size_of(self._shape)
@property
def ndim(self) -> int:
"""Returns the number of dimensions."""
return len(self._shape)
def reshape(self: _DcT, shape: Union[Shape, str], **axes_length: int) -> _DcT:
"""Reshape the batch shape according to the pattern.
Supports both tuple and einops mode:
```python
rays.reshape('b h w -> b (h w)')
rays.reshape((128, -1))
```
Args:
shape: Target shape. Can be string for `einops` support.
**axes_length: Any additional specifications for dimensions for einops
support.
Returns:
The dataclass array with the new shape
"""
if isinstance(shape, str): # Einops support
return self._map_field( # pylint: disable=protected-access
array_fn=lambda f: einops.rearrange( # pylint: disable=g-long-lambda
f.value,
np_utils.to_absolute_einops(shape, nlastdim=len(f.inner_shape)),
**axes_length,
),
dc_fn=lambda f: f.value.reshape( # pylint: disable=g-long-lambda
np_utils.to_absolute_einops(shape, nlastdim=len(f.inner_shape)),
**axes_length,
),
)
else: # Numpy support
assert isinstance(shape, tuple) # For pytest
def _reshape(f: _ArrayField):
return f.value.reshape(shape + f.inner_shape)
return self._map_field(array_fn=_reshape, dc_fn=_reshape) # pylint: disable=protected-access
def flatten(self: _DcT) -> _DcT:
"""Flatten the batch shape."""
return self.reshape((-1,))
def broadcast_to(self: _DcT, shape: Shape) -> _DcT:
"""Broadcast the batch shape."""
return self._map_field( # pylint: disable=protected-access
array_fn=lambda f: f.broadcast_to(shape),
dc_fn=lambda f: f.broadcast_to(shape),
)
def __getitem__(self: _DcT, indices: _IndicesArg) -> _DcT:
"""Slice indexing."""
indices = np.index_exp[indices] # Normalize indices
# Replace `...` by explicit shape
indices = _to_absolute_indices(indices, shape=self.shape)
return self._map_field(
array_fn=lambda f: f.value[indices],
dc_fn=lambda f: f.value[indices], # pyrefly: ignore[bad-argument-type]
)
# _DcT[n *d] -> Iterator[_DcT[*d]]
def __iter__(self: _DcT) -> Iterator[_DcT]:
"""Iterate over the outermost dimension."""
if not self.shape:
raise TypeError(f'iteration over 0-d array: {self!r}')
# Similar to `etree.unzip(self)` (but work with any backend)
field_names = [f.name for f in self._array_fields] # pylint: disable=not-an-iterable
field_values = [f.value for f in self._array_fields] # pylint: disable=not-an-iterable
for vals in zip(*field_values):
yield self.replace(**dict(zip(field_names, vals)))
def __len__(self) -> int:
"""Length of the first array dimension."""
if not self.shape:
raise TypeError(
f'len() of unsized {self.__class__.__name__} (shape={self.shape})'
)
return self.shape[0]
def __bool__(self) -> Literal[True]:
"""`dca.DataclassArray` always evaluate to `True`.
Like all python objects (including dataclasses), `dca.DataclassArray` always
evaluate to `True`. So:
`Ray(pos=None)`, `Ray(pos=0)` all evaluate to `True`.
This allow construct like:
```python
def fn(ray: Optional[dca.Ray] = None):
if ray:
...
```
Or:
```python
def fn(ray: Optional[dca.Ray] = None):
ray = ray or default_ray
```
Only in the very rare case of empty-tensor (`shape=(0, ...)`)
```python
assert ray is not None
assert len(ray) == 0
bool(ray) # TypeError: Truth value is ambigous
```
Returns:
True
Raises:
ValueError: If `len(self) == 0` to avoid ambiguity.
"""
if self.shape and not len(self): # pylint: disable=g-explicit-length-test
raise ValueError(
f'The truth value of {self.__class__.__name__} when `len(x) == 0` '
'is ambigous. Use `len(x)` or `x is not None`.'
)
return True
def map_field(
self: _DcT,
fn: Callable[[Array['*din']], Array['*dout']], # pyrefly: ignore[not-a-type]
) -> _DcT:
"""Apply a transformation on all arrays from the fields."""
return self._map_field( # pylint: disable=protected-access
array_fn=lambda f: fn(f.value),
dc_fn=lambda f: f.value.map_field(fn),
)
# ====== Dataclass/Conversion utils ======
def replace(self: _DcT, **kwargs: Any) -> _DcT:
"""Alias for `dataclasses.replace`."""
init_kwargs = {
k: v for k, v in kwargs.items() if k not in self.__dca_non_init_fields__
}
non_init_kwargs = {
k: v for k, v in kwargs.items() if k in self.__dca_non_init_fields__
}
# Create the new object
new_self = dataclasses.replace(self, **init_kwargs) # pytype: disable=wrong-arg-types # re-none
# TODO(epot): Could try to unify logic bellow with `tree_unflatten`
# Additionally forward the non-init kwargs
# `dataclasses.field(init=False) kwargs are required because `init=True`
# creates conflicts:
# * Inheritance fails with non-default argument 'K' follows default argument
# * Pytype complains too
# TODO(py310): Cleanup using `dataclass(kw_only)`
assert new_self is not self
for k in self.__dca_non_init_fields__:
if k in non_init_kwargs:
v = non_init_kwargs[k]
else:
v = getattr(self, k)
new_self._setattr(k, v) # pylint: disable=protected-access
return new_self
def as_np(self: _DcT) -> _DcT:
"""Returns the instance as containing `np.ndarray`."""
return self.as_xnp(enp.lazy.np)
def as_jax(self: _DcT) -> _DcT:
"""Returns the instance as containing `jnp.ndarray`."""
return self.as_xnp(enp.lazy.jnp)
def as_tf(self: _DcT) -> _DcT:
"""Returns the instance as containing `tf.Tensor`."""
return self.as_xnp(enp.lazy.tnp)
def as_torch(self: _DcT) -> _DcT:
"""Returns the instance as containing `torch.Tensor`."""
return self.as_xnp(enp.lazy.torch)
def as_xnp(self: _DcT, xnp: enp.NpModule) -> _DcT:
"""Returns the instance as containing `xnp.ndarray`."""
if xnp is self.xnp: # No-op
return self
# Direct `torch` <> `tf`/`jax` conversion not supported, so convert to
# `numpy`
if enp.lazy.is_torch_xnp(xnp) or enp.lazy.is_torch_xnp(self.xnp):
def _as_torch(f):
arr = np.asarray(f.value)
# Torch fail for scalar arrays:
# https://github.com/pytorch/pytorch/issues/97021
if enp.lazy.is_torch_xnp(xnp) and not arr.shape: # Destination is torch
return xnp.asarray(arr.item(), dtype=lazy.as_torch_dtype(arr.dtype))
return xnp.asarray(arr)
array_fn = _as_torch
else:
array_fn = lambda f: xnp.asarray(f.value)
# Update all childs
new_self = self._map_field( # pylint: disable=protected-access
array_fn=array_fn,
dc_fn=lambda f: f.value.as_xnp(xnp),
)
return new_self
# TODO(pytype): Remove hack. Currently, Python does not support typing
# annotations for modules, by pytype auto-infer the correct type.
# So this hack allow auto-completion
if typing.TYPE_CHECKING:
@property
def xnp(self): # pylint: disable=function-redefined
"""Returns the numpy module of the class (np, jnp, tnp)."""
return np
else:
@property
def xnp(self) -> enp.NpModule:
"""Returns the numpy module of the class (np, jnp, tnp)."""
return self._xnp
# ====== Torch specific methods ======
# Could also add
# * x.detach
# * x.is_cuda
# * x.device
# * x.get_device
def to(self: _DcT, device, **kwargs) -> _DcT:
"""Move the dataclass array to the device."""
if not lazy.is_torch_xnp(self.xnp):
raise ValueError('`.to` can only be called when `xnp == torch`')
return self.map_field(lambda f: f.to(device, **kwargs))
def cpu(self: _DcT, *args, **kwargs) -> _DcT:
"""Move the dataclass array to the CPU device."""
if not lazy.is_torch_xnp(self.xnp):
raise ValueError('`.cpu` can only be called when `xnp == torch`')
return self.map_field(lambda f: f.cpu(*args, **kwargs))
def cuda(self: _DcT, *args, **kwargs) -> _DcT:
"""Move the dataclass array to the CUDA device."""
if not lazy.is_torch_xnp(self.xnp):
raise ValueError('`.cuda` can only be called when `xnp == torch`')
return self.map_field(lambda f: f.cuda(*args, **kwargs))
# ====== Internal ======
@functools.cached_property
def _all_fields_empty(self) -> bool:
"""Returns True if the `dataclass_array` is invalid."""
if not self._array_fields: # All fields are `None` / `object`
# No fields have been defined.
# This can be the case internally by jax which apply some
# `tree_map(lambda x: sentinel)`.
return True
# `tf.nest` sometimes replace values by dummy `.` inside
# `assert_same_structure`
if enp.lazy.has_tf:
# pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
from tensorflow.python.util import nest_util # pytype: disable=import-error
# pylint: enable=g-direct-tensorflow-import,g-import-not-at-top
if any(f.value is nest_util._DOT for f in self._array_fields): # pylint: disable=protected-access,not-an-iterable
return True
return False
@functools.cached_property
def _all_array_fields(self) -> dict[str, _ArrayField]:
"""All array fields, including `None` values."""
assert self._dca_fields_metadata is not None
return { # pylint: disable=g-complex-comprehension
name: _ArrayField(
name=name,
host=self,
**field_metadata.to_dict(), # pylint: disable=not-a-mapping
)
for name, field_metadata in self._dca_fields_metadata.items() # pylint: disable=protected-access
}
@functools.cached_property
def _array_fields(self) -> list[_ArrayField]:
"""All active array fields (non-None), including static ones."""
# Filter `None` values
return [
f for f in self._all_array_fields.values() if not f.is_value_missing
]
def _cast_xnp_dtype_inplace(self) -> Optional[enp.NpModule]:
"""Validate `xnp` are consistent and cast `np` -> `xnp` in-place."""
if self._all_fields_empty: # pylint: disable=using-constant-test
return None
# Validate the dtype
def _get_xnp(f: _ArrayField) -> enp.NpModule:
try:
return np_utils.get_xnp(
f.value,
strict=not self.__dca_params__.cast_list,
)
except Exception as e: # pylint: disable=broad-except
epy.reraise(e, prefix=f'Invalid {f.qualname}: ')
xnps = epy.groupby(
self._array_fields,
key=_get_xnp,
value=lambda f: f.name,
)
if not xnps:
return None
xnp = _infer_xnp(xnps)
def _cast_field(f: _ArrayField) -> None:
try:
# Supports for TensorSpec (e.g. in `tf.function` signature)
if enp.lazy.is_tf_xnp(xnp) and isinstance(
f.value, enp.lazy.tf.TensorSpec
):
# TODO(epot): Actually check the dtype
new_value = f.value
else:
new_value = np_utils.asarray(
f.value,
xnp=xnp,
dtype=f.dtype, # pyrefly: ignore[bad-argument-type]
cast_dtype=self.__dca_params__.cast_dtype,
)
self._setattr(f.name, new_value)
# After the field has been set, we validate the shape
f.assert_shape()
except Exception as e: # pylint: disable=broad-except
epy.reraise(e, prefix=f'Invalid {f.qualname}: ')
self._map_field(
array_fn=_cast_field,
dc_fn=_cast_field, # pytype: disable=wrong-arg-types
_inplace=True,
)
return xnp
def _broadcast_shape_inplace(self) -> Optional[Shape]:
"""Validate the shapes are consistent and broadcast values in-place."""
if self._all_fields_empty: # pylint: disable=using-constant-test
return None
# First collect all shapes and compute the final shape.
shape_to_names = epy.groupby(
self._array_fields,
key=lambda f: f.host_shape,
value=lambda f: f.name,
)
shape_lengths = {len(s) for s in shape_to_names.keys()}
# Broadcast all shape together
try:
final_shape = np.broadcast_shapes(*shape_to_names.keys())
except ValueError:
final_shape = None # Bad broadcast
# Currently, we restrict broadcasting to either scalar or fixed length.
# This is to avoid confusion broadcasting vs vectorization rules.
# This restriction could be lifted if we encounter a use-case.
if (
final_shape is None
or len(shape_lengths) > 2
or (len(shape_lengths) == 2 and 0 not in shape_lengths)
):
raise ValueError(
f'Conflicting batch shapes: {shape_to_names}. '
f'Currently {type(self).__qualname__}.__init__ broadcasting is '
'restricted to scalar or dim=1 . '
'Please open an issue if you need more fine-grained broadcasting.'
)
def _broadcast_field(f: _ArrayField) -> None:
if f.host_shape == final_shape: # Already broadcasted
return
elif not self.__dca_params__.broadcast: # Broadcasing disabled
raise ValueError(
f'{type(self).__qualname__} has `broadcast=False`. '
f'Cannot broadcast {f.name} from {f.full_shape} to {final_shape}. '
'To enable broadcast, use `@dca.dataclass_array(broadcast=True)`.'
)
self._setattr(f.name, f.broadcast_to(final_shape))
self._map_field(
array_fn=_broadcast_field,
dc_fn=_broadcast_field, # pytype: disable=wrong-arg-types
_inplace=True,
)
return final_shape
def _to_absolute_axis(self, axis: Axes) -> Axes:
"""Normalize the axis to absolute value."""
try:
return np_utils.to_absolute_axis(axis, ndim=self.ndim)
except Exception as e: # pylint: disable=broad-except
epy.reraise(
e,
prefix=f'For {self.__class__.__qualname__} with shape={self.shape}: ',
)
def _map_field(
self: _DcT,
*,
array_fn: Callable[[_ArrayField[Array['*din']]], Array['*dout']], # pyrefly: ignore[not-a-type]
dc_fn: Optional[Callable[[_ArrayField[_DcT]], _DcT]],
_inplace: bool = False,
) -> _DcT:
"""Apply a transformation on all array fields structure.
Args:
array_fn: Function applied on the `xnp.ndarray` fields
dc_fn: Function applied on the `dca.DataclassArray` fields (to recurse)
_inplace: If True, assume the function mutate the object in-place. Should
only be used inside `__init__` for performances.
Returns:
The transformed dataclass array.
"""
def _apply_field_dn(f: _ArrayField):
if f.is_dataclass: # Recurse on dataclasses
return dc_fn(f) # pylint: disable=protected-access # pyrefly: ignore[not-callable]
else:
return array_fn(f)
new_values = {f.name: _apply_field_dn(f) for f in self._array_fields} # pylint: disable=not-an-iterable,protected-access
# For performance, do not call replace to save the constructor call
if not _inplace:
return self.replace(**new_values)
else:
return self
def tree_flatten_with_keys(
self,
) -> tuple[
tuple[tuple[enp.lazy.jax.tree_utils.BuiltInKeyEntry, DcOrArray], ...],
_TreeMetadata,
]:
"""`jax.tree_utils` support."""
tree_util = enp.lazy.jax.tree_util
# We flatten all values (and not just the non-None ones)
array_field_values = tuple(
(tree_util.GetAttrKey(k), f.value)
for k, f in self._all_array_fields.items()
)
metadata = _TreeMetadata(
array_field_names=list(self._all_array_fields.keys()),
non_array_field_kwargs={
f.name: getattr(self, f.name)
for f in dataclasses.fields(self) # pytype: disable=wrong-arg-types # re-none
if f.name not in self._all_array_fields # pylint: disable=unsupported-membership-test
},
)
return (array_field_values, metadata)
@classmethod
def tree_unflatten(
cls: Type[_DcT],
metadata: _TreeMetadata,
array_field_values: list[DcOrArray],
) -> _DcT:
"""`jax.tree_utils` support."""
return cls._tree_unflatten(
metadata=metadata,
array_field_values=array_field_values,
constructor=cls,
)
@classmethod
def _tree_unflatten(
cls: Type[_DcT],
*,
metadata: _TreeMetadata,
array_field_values: list[DcOrArray],
constructor: Callable[..., _DcT],
) -> _DcT:
"""Initialize a model after serialization."""
array_field_kwargs = dict(
zip(
metadata.array_field_names,
array_field_values,
)
)
init_fields = {}
non_init_fields = {}
fields = {f.name: f for f in dataclasses.fields(cls)} # pytype: disable=wrong-arg-types # re-none
for k, v in metadata.non_array_field_kwargs.items():
if fields[k].init:
init_fields[k] = v
else:
non_init_fields[k] = v
self = constructor(**array_field_kwargs, **init_fields)
# Currently it's not clear how to handle non-init fields so raise an error
if non_init_fields:
if set(non_init_fields) - self.__dca_non_init_fields__:
raise ValueError(
'`dca.DataclassArray` field with init=False should be explicitly '
'specified in `__dca_non_init_fields__` for them to be '
'propagated by `tree_map`.'
)
# TODO(py310): Delete once dataclass supports `kw_only=True`
for k, v in non_init_fields.items():
self._setattr(k, v) # pylint: disable=protected-access
return self
def _tree_flatten(self) -> tuple[tuple[DcOrArray, ...], _TreeMetadata]:
components, metadata = self.tree_flatten_with_keys()
components = tuple(v for _, v in components)
return components, metadata
def __tf_flatten__(self) -> tuple[_TreeMetadata, tuple[DcOrArray, ...]]:
components, metadata = self._tree_flatten()
return metadata, components
@classmethod
def __tf_unflatten__(
cls: Type[_DcT],
metadata: _TreeMetadata,
components: list[DcOrArray],
) -> _DcT:
return cls.tree_unflatten(metadata, components)
def __getstate__(self) -> tuple[_TreeMetadata, tuple[DcOrArray, ...]]:
components, metadata = self._tree_flatten()
return metadata, components
def __setstate__(self, state: tuple[_TreeMetadata, list[DcOrArray]]) -> None:
def constructor(**kwargs):
self.__init__(**kwargs)
return self
metadata, components = state
type(self)._tree_unflatten(
metadata=metadata,
array_field_values=components,
constructor=constructor,
)
def _setattr(self, name: str, value: Any) -> None:
"""Like setattr, but support `frozen` dataclasses."""
object.__setattr__(self, name, value)
def assert_same_xnp(self, x: Union[Array[...], DataclassArray]) -> None: # pyrefly: ignore[not-a-type]
"""Assert the given array is of the same type as the current object."""
xnp = np_utils.get_xnp(x)
if xnp is not self.xnp:
raise ValueError(
f'{self.__class__.__name__} is {self.xnp.__name__} but got input '
f'{xnp.__name__}. Please cast input first.'
)
def _init_cls(self: DataclassArray) -> None:
"""Setup the class the first time the instance is called.
This will:
* Extract the types annotations, detect which fields are arrays or static,
and store the result in `_dca_fields_metadata`
* For static `DataclassArray` (class with only static fields), it will
add a dummy array field for compatibility with `.xnp`/`.shape` (so
methods works correctly and return the right shape/xnp when nested)
Args:
self: The dataclass to initialize
"""
cls = type(self)
# The first time, compute typing annotations & metadata
# At this point, `ForwardRef` should have been resolved.
try:
hints = typing_extensions.get_type_hints(cls, include_extras=True)
except Exception as e: # pylint: disable=broad-except
msg = (
f'Could not infer typing annotation of {cls.__qualname__} '
f'defined in {cls.__module__}:\n'
)
lines = [f' * {k}: {v!r}' for k, v in cls.__annotations__.items()]
lines = '\n'.join(lines)
epy.reraise(e, prefix=msg + lines + '\n')
# TODO(epot): Remove restriction once pytype supports `datclass_transform`
# and `dca` automatically apply the `@dataclasses.dataclass`
if _DUMMY_ARRAY_FIELD in cls.__dataclass_fields__: # pytype: disable=attribute-error
raise NotImplementedError(
'Suclassing of DataclassArray with no array field is not supported '
'after an instance of the class was created. Error raised for '
f'{cls.__qualname__}'
)
dca_fields_metadata = {
f.name: _make_field_metadata(f, hints) for f in dataclasses.fields(cls) # pytype: disable=wrong-arg-types
}
dca_fields_metadata = { # Filter `None` values (static fields)
k: v for k, v in dca_fields_metadata.items() if v is not None
}
if not dca_fields_metadata:
# DataclassArray without any array fields
# Hack: To support `.xnp`, `.shape`, we add a dummy empty field which
# is propagated by the various ops.
dca_fields_metadata[_DUMMY_ARRAY_FIELD] = _ArrayFieldMetadata( # pytype: disable=wrong-arg-types
inner_shape_non_static=(),
dtype=np.float32, # pyrefly: ignore[bad-argument-type]
)
default_dummy_array = np.zeros((), dtype=np.float32)
_add_field_to_dataclass(
cls, _DUMMY_ARRAY_FIELD, default=default_dummy_array
)
# Because we're in `__init__`, so also update the current call
self._setattr(_DUMMY_ARRAY_FIELD, default_dummy_array) # pylint: disable=protected-access
cls._dca_fields_metadata = ( # pylint: disable=protected-access
dca_fields_metadata
)
def _add_field_to_dataclass(cls, name: str, default: Any) -> None:
"""Add a new field to the given dataclass."""
# Make sure to not update the parent class
# Otherwise we could even accidentally update `dca.DataclassArray`
if '__dataclass_fields__' not in cls.__dict__:
# TODO(epot): Remove the limitation once `dataclasses.dataclass` is
# automatically applied
raise ValueError(
f'{cls.__name__} is not a `@dataclasses.dataclass(frozen=True)`'
)
assert name not in cls.__dataclass_fields__ # pytype: disable=attribute-error
# Ideally, we want init=False, so sub-dataclass ignore this field
# but this makes `.replace` fail
field = dataclasses.field(default=default, init=True, repr=False)
field.__set_name__(cls, name)
field.name = name
field.type = Any
field._field_type = dataclasses._FIELD # pylint: disable=protected-access # pytype: disable=module-attr
cls.__dataclass_fields__[name] = field # pytype: disable=attribute-error
original_init = cls.__init__
@functools.wraps(original_init)
def new_init(self, **kwargs: Any):
self._setattr(name, kwargs.pop(name, default)) # pylint: disable=protected-access
return original_init(self, **kwargs)
cls.__init__ = new_init
def _infer_xnp(xnps: dict[enp.NpModule, list[str]]) -> enp.NpModule:
"""Extract the `xnp` module."""
non_np_xnps = set(xnps) - {np} # jnp, tnp take precedence on `np`
# Detecting conflicting xnp
if len(non_np_xnps) > 1:
xnps = {k.__name__: v for k, v in xnps.items()}
raise ValueError(f'Conflicting numpy types: {xnps}')
if not non_np_xnps:
return np