class_docstr: Estimator class based on scikit-learn estimators.
  
  See https://scikit-learn.org/stable/developers/develop.html
class_docstr_tok: ["Estimator","class","based","on","scikit-learn","estimators",".","See","https",":","\/\/scikit-learn.org\/stable\/developers\/develop.html"]
docstr: Predict quantiles over the time points in X.
  
  This uses the fit params_ in this class, which has a single batch dimension.
  By default, the function will compute a quantiles over this batch dimension,
  but can compute per-particle quantiles by passing `axis=()`.
  
  Args:
    X: Training array, where the rows are samples, and the column are the
      number of features (assumed to be 1).
    q: Quantiles in [0, 100] to compute.
    axis: Tuple of dimensions to reduce over.
  
  Raises:
    RuntimeError: If model is not fit.
  
  Returns:
    Array with a row for each element of `q`.
docstr_tok: ["Predict","quantiles","over","the","time","points","in","X",".","This","uses","the","fit","params_","in","this","class",",","which","has","a","single","batch","dimension",".","By","default",",","the","function","will","compute","a","quantiles","over","this","batch","dimension",",","but","can","compute","per-particle","quantiles","by","passing","`","axis=","(",")","`",".","Args",":","X",":","Training","array",",","where","the","rows","are","samples",",","and","the","column","are","the","number","of","features","(","assumed","to","be","1",")",".","q",":","Quantiles","in","[","0",",","100","]","to","compute",".","axis",":","Tuple","of","dimensions","to","reduce","over",".","Raises",":","RuntimeError",":","If","model","is","not","fit",".","Returns",":","Array","with","a","row","for","each","element","of","`","q","`","."]
code: def predict_quantiles(
      self,
      X: jax.Array,
      q=(2.5, 50.0, 97.5),
      axis: tuple[int, ...] = (0,),  # pylint: disable=invalid-name
  ) -> jax.Array:
      """Predict quantiles over the time points in X.
  
      This uses the fit params_ in this class, which has a single batch dimension.
      By default, the function will compute a quantiles over this batch dimension,
      but can compute per-particle quantiles by passing `axis=()`.
  
      Args:
        X: Training array, where the rows are samples, and the column are the
          number of features (assumed to be 1).
        q: Quantiles in [0, 100] to compute.
        axis: Tuple of dimensions to reduce over.
  
      Raises:
        RuntimeError: If model is not fit.
  
      Returns:
        Array with a row for each element of `q`.
      """
      if not self.check_is_fitted():
          raise RuntimeError(
              "Model is not yet fit! Call `AutoBNN.fit(X, y)` first."
          )
      loc = training_util.make_predictions(
          self.params_, net=self.net_, x_test=X
      )
      assert hasattr(self.likelihood_, "sample")
  
      # We are doing all quantiles via sampling currently, so the fact that we
      # also reduce over the new batch axis is an implementation detail.
      axis = (0,) + tuple(j + 1 for j in axis)
      draws = self.likelihood_.sample(
          self.params_["params"],
          jnp.squeeze(loc),
          seed=self.predict_quantiles_seed_,
          sample_shape=100,
      )
      # TODO(ursk): return these in a dictionary rather than as a tuple
      return jnp.percentile(draws, jnp.array(q), axis=axis)
  
imports: from typing import Any, Mapping, Optional, Sequence
  import jax
  import jax.numpy
  from jaxtyping import ArrayLike, PyTree
  from tensorflow_probability.python.experimental.autobnn import bnn
  from tensorflow_probability.python.experimental.autobnn import likelihoods
  from tensorflow_probability.python.experimental.autobnn import models
  from tensorflow_probability.python.experimental.autobnn import training_util
label_desc: Use image node_id 8 for calling the _AutoBnnEstimator obj's underlying member method code with example usage: obj.predict_quantiles(X, q, axis) and returns: jnp