Currently there is no way of recovering the inner spec.
This is an issue for example when using nested datasets when `ds.element_spec` will return the outer dataset spec `{'nested_ds': DatasetSpec({'img': TensorSpec(...)})})` but it's not possible to access the inner DatasetSpec.
PiperOrigin-RevId: 353636988
Change-Id: I4bcfb3ab31a0761834a2837075264f8117973861
This is the right behavior for TF and JAX and matches what TF does on GPU. It
doesn't match TF on CPU, but that's really a TF bug.
PiperOrigin-RevId: 353628258
Change-Id: Iadb0ce627e033b3fd730e899501f5818d6ecc78d
This is the right behavior for TF and JAX and matches what TF does on GPU. It
doesn't match TF on CPU, but that's really a TF bug.
PiperOrigin-RevId: 353624935
Change-Id: Ic54e718a45e412db031fd1b3772d135ac46ed8a1
Allow native segment execution for BinaryTensorWeightBroadcastTest when
TensorRT 7+ is used. This is to workaround b/176540862.
PiperOrigin-RevId: 353581801
Change-Id: Ic494e212c012a3fb2c260da64a191899e68ec866
This is required for enabling flex fallback for complicated variable operators
for stateful random operators and so on.
PiperOrigin-RevId: 353578616
Change-Id: Ic71f36199ff437c2efa6f259ed2fd474694617d9
2. Ensure the specified input dimension size is same w/ that specified in the model.
PiperOrigin-RevId: 353569125
Change-Id: I3df81a9314b97689e1925ab88fc2e575f9579ffa
Added pooling_test_util with unified tests.
Added Metal pooling unit tests.
PiperOrigin-RevId: 353513964
Change-Id: I5a58b319509b64f268904312611e4649b35a467a
Added max_unpooling_test_util with unified tests.
Added Metal max unpooling unit test.
PiperOrigin-RevId: 353509117
Change-Id: If9a658f8be8786e731a4fb72ede6b2f023b601a9
Now Metal Mean supports any combination of axis reduction.
Metal supports reduction ops with any combination of axis.
PiperOrigin-RevId: 353507802
Change-Id: Ib032b83e6075c1004e5249b6b8c8f66a13b2cf79
These are not actionable by our end users, and it usually just confuses them.
PiperOrigin-RevId: 353372479
Change-Id: I0994bfd35f0a5145349f54e26b35b11bcb85f0ed
induction variables (AIV). Specifically, candidates are gtes, e.g.,
gte(param0, N). We check if the loop body plumbs the AIV
through the same tuple index at root, and that ops involving AIV
involve constants.
op2 = op(constants, gte(param0, N), constants)
op3 = op(constants, f(op2, gte(param0, N), constants)
root = tuple(..., op3, ...)
Further, the ops are restricted to basic math ops (+,-,*,/).
PiperOrigin-RevId: 353357828
Change-Id: If2bbd7a4d4758982100c2a29e618ac40abf1de1e
This is no longer needed as tpu_extract_outside_compilation correctly moves ops to host and minimizes communication without explicit clustering.
PiperOrigin-RevId: 353357796
Change-Id: I5babc694680991e5b7dcbbe4326180adeeb22c8e
Also remove the use of async executor to launch collective ops in eager mode and use one thread per device instead. This resolves the issue of not being able to call numpy() on the result of async executor. This change applies to MWMS too.
PiperOrigin-RevId: 353355403
Change-Id: I9c9f30dfe18dc830a4a8fa9bbaec042c7c2edd8f
