jax-js
    Preparing search index...

    Function valueAndGrad

    • Create a function that evaluates both f and the gradient of f.

      When { hasAux: true } is passed, the function f is expected to return an [out, aux] tuple, and the return value will be [[out, aux], gradient].

      Type Parameters

      • F extends (...args: any[]) => JsTree<Array>
      • const I extends number | number[] | undefined = undefined
      • const HA extends boolean = false

      Parameters

      • f: F
      • Optionalopts: Omit<{ argnums?: number | number[]; hasAux?: boolean }, "argnums"> & {
            argnums?: I;
            hasAux?: HA;
        }

      Returns (
          ...primals: MappedJsTree<Parameters<F>, Array, ArrayLike>,
      ) => [
          ReturnType<F>,
          MappedJsTree<
              I extends undefined
                  ? Parameters<F>[0]
                  : I extends number
                      ? Parameters<F>[I<I>]
                      : I extends number[]
                          ? {
                              [K in string
                              | number
                              | symbol]: I<I>[K] extends number ? Parameters<F>[any[any]] : never
                          }
                          : never,
              ArrayLike,
              Array,
          >,
      ]

      // Without hasAux
      const [value, gradient] = valueAndGrad(f)(x);

      // With hasAux
      const [[value, aux], gradient] = valueAndGrad(f, { hasAux: true })(x);