Enhance the Tensor-Tracer in the following ways:
(1) Able to trace tensors when the model is executed on the CPU.
(previously, it could only trace when the model is executed on TPU)
(2) Allow the user to specify the op-names and op-types to be excluded or included for tracing via regular expressions.
(3) Two new trace modes: (1) tracing the vector norm of the tensor and (2) tracing the maximum of the absolute values of all elements in the tensor.
(4) Attach the replica-ID to a traced tensor value so that the post-processing tool (Tensor-Inspector) can reconstruct the whole tensor from all replicas.
(5) An API to trace tensors programmatically.
(6) Allow writing the trace to stdout (previously, it must be written to a file).
PiperOrigin-RevId: 225112219
Loading
Please sign in to comment