[WIP] Amend documentation of JAX config options, include context managers, mechanism to hide flags #28006
Conversation
| capitalization_rules_list = [ | ||
| ("jax", "JAX"), ("xla", "XLA"), ("pgle", "PGLE"), ("cuda", "CUDA"), ("vjp", "VJP"), ("jvp", "JVP"), | ||
| ("pjrt", "PjRT"), ("gpu", "GPU"), ("tpu", "TPU"), ("prng", "PRNG"), ("roocm", "ROOCM"), ("spmd", "SPMD"), | ||
| ("bcoo", "BCOO"), ("jit", "JIT"), ("cpu", "CPU"), ("cusparse", "cuSPARSE"), ("ir", "IR"), ("dtype", "DType"), |
There was a problem hiding this comment.
jit should be lowercaps.
roocm is misspelled? rocm
There was a problem hiding this comment.
Thanks, fixed rocm. These capitalization rules just get used to build the titles in the doc, so as an acronym jit I'm leaving jit capitalised.
There was a problem hiding this comment.
PjRT -> PJRT
What about jit -> jit? Since it probably refers to the JAX API, not "Just In Time" (jax's API is kind of a misuse of the generic term, the verb is really "JIT compile")
Also numpy uses dtype, not Dtype, even when it's not dtype. There are examples in https://numpy.org/doc/2.1/reference/arrays.dtypes.html.
There was a problem hiding this comment.
Oh I see Dtype because it's all Title Case. This is really yak shaving now but I personally prefer just capitalizing the first word, I think they look funny as Full Titles
Also nan -> NaN
No description provided.