Module Einsum_parser

Entry point for the einsum parser library.

This module provides functions to parse einsum notation specifications using a Menhir-based parser.

type use_padding_spec = [
  1. | `True
  2. | `False
  3. | `Unspecified
]

Use_padding specification for convolutions.

val sexp_of_use_padding_spec : use_padding_spec -> Sexplib0.Sexp.t
type conv_spec = {
  1. dilation : Base.string;
  2. kernel_label : Base.string;
  3. use_padding : use_padding_spec;
}

Convolution component for affine axis specifications. Note: dilation is a string because it can be an identifier at parse time, and is resolved to an int at runtime.

val sexp_of_conv_spec : conv_spec -> Sexplib0.Sexp.t
type axis_spec =
  1. | Label of Base.string
    (*

    A variable axis label.

    *)
  2. | Fixed_index of Base.int
    (*

    A fixed index, used for projection.

    *)
  3. | Affine_spec of {
    1. stride : Base.string;
      (*

      Coefficient for the over dimension (string to allow identifiers).

      *)
    2. over_label : Base.string;
      (*

      The output/iteration dimension label.

      *)
    3. conv : conv_spec Base.option;
      (*

      Optional convolution: dilation*kernel.

      *)
    4. stride_offset : Base.int;
      (*

      Constant offset added after stride*over.

      *)
    }
    (*

    Affine axis specification: stride*over + stride_offset + dilation*kernel. Corresponds to Row.Affine in shape inference.

    *)
  4. | Concat_spec of concat_component Base.list
    (*

    Concatenation of several components into a single axis. Corresponds to Row.Concat in shape inference.

    *)

Specification for individual axes in the einsum notation. Note: stride is a string because it can be an identifier at parse time, and is resolved to an int at runtime.

and concat_component =
  1. | Concat_label of Base.string
    (*

    A dimension variable, shared with the spec's other slots.

    *)
  2. | Concat_size of Base.int
    (*

    A component of a fixed size that no other slot names: in 3^a => a it is skipped, in a => a^3 it is left to the result's initialization (gh-ocannl-1656).

    *)

A component of a concatenated axis (a^b, 3^a^5).

val sexp_of_axis_spec : axis_spec -> Sexplib0.Sexp.t
val sexp_of_concat_component : concat_component -> Sexplib0.Sexp.t
module AxisKey : sig ... end

An index pointing to any of a shape's axes, including the kind of the axis (Batch, Input, Output) and the position (which is counted from the end to facilitate broadcasting).

type 'a axis_map = 'a Base.Map.M(AxisKey).t
val sexp_of_axis_map : 'a. ('a -> Sexplib0.Sexp.t) -> 'a axis_map -> Sexplib0.Sexp.t
type parsed_axis_labels = {
  1. bcast_batch : Base.string Base.option;
  2. bcast_input : Base.string Base.option;
  3. bcast_output : Base.string Base.option;
  4. implicit_batch : Base.bool;
  5. implicit_input : Base.bool;
  6. given_batch : axis_spec Base.list;
  7. given_input : axis_spec Base.list;
  8. given_output : axis_spec Base.list;
  9. given_beg_batch : axis_spec Base.list;
  10. given_beg_input : axis_spec Base.list;
  11. given_beg_output : axis_spec Base.list;
  12. labels : axis_spec axis_map;
}

The labels are strings assigned to AxisKey axes. Moreover the bcast_ fields represent whether additional leading/middle axes are allowed (corresponding to the dot-ellipsis syntax for broadcasting). The string can be used to identify a row variable; parsing "..." yields the reserved per-kind names "...batch", "...input", "...output", which cannot collide with user-written identifiers. The implicit_ fields record that the row was omitted together with its kind separator, in which case it reads as the context ellipsis but its row variable may be silently closed to an empty row (even on parameters) when nothing else constrains it — unlike an explicitly written ellipsis. The given_ fields are lists of axis specs of the corresponding kind in labels where from_end=true, given_beg_ where from_end=false.

val sexp_of_parsed_axis_labels : parsed_axis_labels -> Sexplib0.Sexp.t
exception Parse_error of Base.string

Exception raised when parsing fails.

val binary_operators_with_generated_specs : Base.string Base.list
val unary_operators_with_generated_specs : Base.string Base.list
val concat_operator_with_generated_specs : Base.string
val operators_with_generated_specs : Base.string Base.list

The unqualified DSL operators whose literal specifications the OCANNL PPXs pass through identifier substitution. The PPX dispatch tables and source-policy scans share these values so adding a syntax form cannot silently update only one consumer.

val is_multichar : Base.string -> Base.bool

Determine if a spec uses multichar mode. Multichar mode is triggered by presence of: ',', '*', '+', '^', '&', '<', or a '=' that does not begin "=>". In multichar mode axes are separated by commas; whitespace does not separate them.

val axis_labels_of_spec : Base.string -> parsed_axis_labels

Parse an axis labels specification.

Examples:

  • "abc" (single-char mode)
  • "a, b, c" (multichar mode, triggered by comma)
  • "batch|input->output"
  • "...a..b"
val einsum_of_spec : Base.string -> parsed_axis_labels Base.list * parsed_axis_labels

Parse an einsum specification. Returns a list of RHS specs and the LHS (result) spec.

Examples:

  • "ij;jk=>ik" (matrix multiplication, 2 RHSes)
  • "ij=>ji" (transpose/permute, 1 RHS)
  • "a;b;c=>d" (3 RHSes for block tensors/concatenation)
  • "i,j=>2*i+j" (convolution)