@thi.ng/tensors/defoprt.js

0D/1D/2D/3D/4D tensors with extensible polymorphic operations and customizable storage

import { identity } from "@thi.ng/api/fn";
import { top } from "./top.js";
const defOpRT = (rfn, init, complete = identity) => {
  const f0 = (a) => complete(rfn(init(), a.data, a.offset), a);
  const f1 = (a) => {
    const {
      data,
      offset,
      shape: [sx],
      stride: [tx]
    } = a;
    let res = init();
    for (let x = 0; x < sx; x++) {
      res = rfn(res, data, offset + x * tx);
    }
    return complete(res, a);
  };
  const f2 = (a) => {
    const {
      data,
      offset,
      shape: [sx, sy],
      stride: [tx, ty]
    } = a;
    let res = init();
    let ox;
    for (let x = 0; x < sx; x++) {
      ox = offset + x * tx;
      for (let y = 0; y < sy; y++) {
        res = rfn(res, data, ox + y * ty);
      }
    }
    return complete(res, a);
  };
  const f3 = (a) => {
    const {
      data,
      offset,
      shape: [sx, sy, sz],
      stride: [tx, ty, tz]
    } = a;
    let res = init();
    let ox, oy;
    for (let x = 0; x < sx; x++) {
      ox = offset + x * tx;
      for (let y = 0; y < sy; y++) {
        oy = ox + y * ty;
        for (let z = 0; z < sz; z++) {
          res = rfn(res, data, oy + z * tz);
        }
      }
    }
    return complete(res, a);
  };
  const f4 = (a) => {
    const {
      data,
      offset,
      shape: [sx, sy, sz, sw],
      stride: [tx, ty, tz, tw]
    } = a;
    let res = init();
    let ox, oy, oz;
    for (let x = 0; x < sx; x++) {
      ox = offset + x * tx;
      for (let y = 0; y < sy; y++) {
        oy = ox + y * ty;
        for (let z = 0; z < sz; z++) {
          oz = oy + z * tz;
          for (let w = 0; w < sw; w++) {
            res = rfn(res, data, oz + w * tw);
          }
        }
      }
    }
    return complete(res, a);
  };
  return top(
    0,
    f0,
    f1,
    f2,
    f3,
    f4
  );
};
export {
  defOpRT
};