Sergei Lebedev daa81e6fb5 Added support for printing scalar values in Pallas TPU kernels ...

The implementation uses the new tpu.log operation in the Mosaic TPU dialect.

Note that

* the logging only happens if --xla_tpu_enable_log_recorder is set;
* only scalars can be printed;
* placeholders only accept i32 arguments at the moment.

PiperOrigin-RevId: 636585852

2024-05-23 10:02:00 -07:00

..

mosaic

Added support for printing scalar values in Pallas TPU kernels

2024-05-23 10:02:00 -07:00

mosaic_gpu

Added support for printing scalar values in Pallas TPU kernels

2024-05-23 10:02:00 -07:00

triton

Added pl.debug_print() -- a new primitive for printing from Pallas kernels

2024-05-22 04:41:42 -07:00

__init__.py

[Pallas] Upstream pallas to JAX

2023-08-01 16:43:13 -07:00

BUILD

Fix some load orderings for buildifier

2024-03-27 10:28:57 -07:00

core.py

Add dynamic grid support to emit_pipeline

2024-04-09 23:55:51 -07:00

pallas_call.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

primitives.py

Added support for printing scalar values in Pallas TPU kernels

2024-05-23 10:02:00 -07:00

utils.py

[Pallas TPU] Convert pattern_match_while_to_fori_loop to return (Jaxpr, str) rather than throw exceptions.

2024-04-09 14:04:25 -07:00