Package: dataclass-array
Dataclass ArrayDataclassArrayare dataclasses which behave like numpy-like arrays (can be
batched, reshaped, sliced,...), compatible with Jax, TensorFlow, and numpy (with
torch support planned).This reduce boilerplate and improve readability. See themotivating examplessection bellow.To view an example of dataclass arrays used in practice, seevisu3d.DocumentationDefinitionTo create adca.DataclassArray, take a frozen dataclass and:Inherit fromdca.DataclassArrayAnnotate the fields withdataclass_array.typingto specify the inner shape
and dtype of the array (see below for static or nested dataclass fields).
The array types are an alias frometils.array_types.importdataclass_arrayasdcafromdataclass_array.typingimportFloatArrayclassRay(dca.DataclassArray):pos:FloatArray['*batch_shape 3']dir:FloatArray['*batch_shape 3']UsageAfterwards, the dataclass can be used as a numpy array:ray=Ray(pos=jnp.zeros((3,3)),dir=jnp.eye(3))ray.shape==(3,)# 3 rays batched togetherray.pos.shape==(3,3)# Individual fields still available# Numpy slicing/indexing/maskingray=ray[...,1:2]ray=ray[norm(ray.dir)>1e-7]# Shape transformationray=ray.reshape((1,3))ray=ray.reshape('h w -> w h')# Native einops supportray=ray.flatten()# Stack multiple dataclass arrays togetherray=dca.stack([ray0,ray1,...])# Supports TF, Jax, Numpy (torch planned) and can be easily convertedray=ray.as_jax()# as_np(), as_tf()ray.xnp==jax.numpy# `numpy`, `jax.numpy`, `tf.experimental.numpy`# Compatibility `with jax.tree_util`, `jax.vmap`,..ray=jax.tree_util.tree_map(lambdax:x+1,ray)ADataclassArrayhas 2 types of fields:Array fields: Fields batched like numpy arrays, with reshape, slicing,...
Can bexnp.ndarrayor nesteddca.DataclassArray.Static fields: Other non-numpy field. Are not modified by reshaping,...
Static fields are also ignored injax.tree_map.classMyArray(dca.DataclassArray):# Array fieldsa:FloatArray['*batch_shape 3']# Defined by `etils.array_types`b:FloatArray['*batch_shape _ _']# Dynamic shapec:Ray# Nested DataclassArray (equivalent to `Ray['*batch_shape']`)d:Ray['*batch_shape 6']# Array fields explicitly definede:Any=dca.field(shape=(3,),dtype=np.float32)f:Any=dca.field(shape=(None,None),dtype=np.float32)# Dynamic shapeg:Ray=dca.field(shape=(3,),dtype=Ray)# Nested DataclassArray# Static field (everything not defined as above)static0:floatstatic1:np.arrayVectorization@dca.vectorize_methodallow your dataclass method to automatically support
batching:Implement method as ifself.shape == ()Decorate the method withdca.vectorize_methodclassCamera(dca.DataclassArray):K:FloatArray['*batch_shape 4 4']resolution=tuple[int,int]@dca.vectorize_methoddefrays(self)->Ray:# Inside `@dca.vectorize_method` shape is always guarantee to be `()`assertself.shape==()assertself.K.shape==(4,4)# Compute the ray as if there was only a single camerareturnRay(pos=...,dir=...)Afterward, we can generate rays for multiple camera batched together:cams=Camera(K=K)# K.shape == (num_cams, 4, 4)rays=cams.rays()# Generate the rays for all the camerascams.shape==(num_cams,)rays.shape==(num_cams,h,w)@dca.vectorize_methodis similar tojax.vmapbut:Only work ondca.DataclassArraymethodsInstead of vectorizing a single axis,@dca.vectorize_methodwill vectorize
over*self.shape(not justself.shape[0]). This is like ifvmapwas
applied toself.flatten()When multiple arguments, axis with dimension1are broadcasted.For example, with__matmul__(self, x: T) -> T:()@(*x,)->(*x,)(b,)@(b,*x)->(b,*x)(b,)@(1,*x)->(b,*x)(1,)@(b,*x)->(b,*x)(b,h,w)@(b,h,w,*x)->(b,h,w,*x)(1,h,w)@(b,1,1,*x)->(b,h,w,*x)(a,*x)@(b,*x)->Error:Incompatiblea!=bTo test on Colab, see thevisu3ddataclassColab tutorial.Motivating examplesdca.DataclassArrayimprove readability by simplifying common patterns:Reshaping all fields of a dataclass:Before (raysis simpledataclass):num_rays=math.prod(rays.origins.shape[:-1])rays=jax.tree_map(lambdar:r.reshape((num_rays,-1)),rays)After (raysisDataclassArray):rays=rays.flatten()# (b, h, w) -> (b*h*w,)Rendering a video:Before (cams: list[Camera]):img=cams[0].render(scene)imgs=np.stack([cam.render(scene)forcamincams[::2]])imgs=np.stack([cam.render(scene)forcamincams])After (cams: Camerawithcams.shape == (num_cams,)):img=cams[0].render(scene)# Render only the first camera (to debug)imgs=cams[::2].render(scene)# Render 1/2 frames (for quicker iteration)imgs=cams.render(scene)# Render all cameras at onceInstallationpipinstalldataclass_arrayThis is not an official Google product