Most docstrings for jax.lax APIs are just short stubs at the moment. We should add documentation including:
- one-line short description
- what
primitive & HLO operation they lower to
Args: section mentioning shape/dtype requirements (if any)Returns: section mentioning output shape/dtypeNotes: section with unexpected conventions or dtype-dependent behaviorExamples: section, but only when they would add clarity (e.g. we don't need an example showing that lax.sin computes the sine)
#26220 is an example of updating a few docstrings.
Most docstrings for
jax.laxAPIs are just short stubs at the moment. We should add documentation including:primitive &HLO operation they lower toArgs:section mentioning shape/dtype requirements (if any)Returns:section mentioning output shape/dtypeNotes:section with unexpected conventions or dtype-dependent behaviorExamples:section, but only when they would add clarity (e.g. we don't need an example showing thatlax.sincomputes the sine)#26220 is an example of updating a few docstrings.