tree 0c8d2983c2f836c33baf50fec50624391e878a20
parent 36caeaec43a80c7c1d9219e7448811a36923798d
author Jean-Baptiste Lespiau <jblespiau@google.com> 1605874928 -0800
committer TensorFlower Gardener <gardener@tensorflow.org> 1605875396 -0800

Add methods to interact with `DeviceArray` objects.

One internal breaking change is subtle:
- `DeviceArray(...).device_buffer.shape()` was a method returning a C++ Shape, now shape is a property, that returns a tuple. In practice, this is only some internal JAX details that no one should use, but I want to explicit this change. shape() has been renamed to xla_shape() which return the C++ object (the Python code has already been uptated to use xla_shape if available).

PiperOrigin-RevId: 343470035
Change-Id: Icea2a6b80c22f9b68e2972efd9945a016171ead2