jax
3b386790 - Summarize large constants that may appear in Jaxpr literals

Commit
1 year ago
Summarize large constants that may appear in Jaxpr literals Once we enable `config.use_simplified_jaxpr_constants` we can end up with potentially large constants in Jaxprs. Instead of printing non-scalars, we summarize them as `[...]`. For example, if before a Jaxpr with constants would be printed as: ``` { lambda a:f32[256]; b:f32[]. let in (b, 1.0:f32[], a) } ``` now would be printed as ``` { lambda a:f32[]. let in (a, 1.0:f32[], [...]:f32[256) } ``` In particular, this summarization does not increse the number of lines in Jaxpr printing.
Author
Committer
Parents
Loading