import { gatherSub as j } from "../ops/gatherSub.js";
import { scatterSub as T } from "../ops/scatterSub.js";
import { t as H, a1 as G, h as A, m as x, d as U, a2 as D, b as q, a3 as F } from "../index-DSGwv2Yx.js";
import { m as J, l as K, a as E, n as N, w as O } from "../not_equal-CkQKkKZy.js";
import { s as m } from "../sum-D39FeU5h.js";
import { s as Q } from "../softmax-BLGJqdwx.js";
function R(e, h, t, g, v, a = 0) {
  return H(() => {
    const S = e.shape[e.shape.length - 1], i = v || e.shape.slice(0, -1), l = i.reduce((n, p) => n * p, 1), f = e.shape.length > 2 ? e.reshape([l, S]) : e, b = h.shape.length > 1 ? h.reshape([l]).cast("int32") : h.cast("int32"), u = J(f, -1, !0), o = G(f, u), r = K(o, -1), d = j(r, b, o);
    let s;
    if (a > 0) {
      const n = G(r, E(o, -1));
      s = A(x(d, 1 - a), x(n, a));
    } else
      s = d;
    if (t)
      if (s = x(s, t), g) {
        const n = m(t.reshape(i), -1);
        s = U(m(s.reshape(i), -1), n);
      } else {
        const n = m(t);
        s = U(m(s), n);
      }
    else
      g ? s = E(s.reshape(i), -1) : s = E(s);
    return s;
  });
}
function ts(e, h, t = 0) {
  return D(
    // @ts-expect-error Invalid params
    (a, S, i) => {
      const l = a.shape[a.shape.length - 1], f = a.shape.slice(0, -1), b = f.reduce((p, c) => p * c, 1), u = a.reshape([b, l]), o = S.reshape([b]).cast("int32");
      let r, d = null;
      if (e) {
        const p = q(-100, "int32"), c = N(o, p);
        d = c.cast("float32"), r = O(c, o, F(o));
      } else
        r = o;
      const s = R(
        u,
        r,
        d || void 0,
        h,
        f,
        t
      );
      return i(d ? [u, r, d] : [u, r]), u.dispose(), o.dispose(), { value: s, gradFunc: (p, c) => H(() => {
        const C = c[0], z = c[1], y = e ? c[2] : void 0, w = Q(C), P = y ? m(y) : q(C.shape[0], "float32"), I = p.div(P).broadcastTo([C.shape[0]]), M = y && e ? x(I, y) : I;
        let L;
        if (t > 0) {
          const W = t / l, $ = G(w, W);
          L = T(
            $,
            z,
            x(M, 1 - t)
          );
        } else
          L = T(w, z, M);
        const V = F(S);
        return [L.reshape(a.shape), V];
      }) };
    }
  );
}
export {
  ts as createSoftmaxCrossEntropyWithGrad,
  R as sparseSoftmaxCrossEntropy
};
