kumi 0.0.15 → 0.0.17

data/lib/kumi/core/functions/registry_v2.rb

@@ -1,68 +0,0 @@
-module Kumi::Core::Functions
-  class RegistryV2
-    def initialize(functions:)
-      @by_name = functions.group_by(&:name).transform_values { |v| v.sort_by(&:opset).freeze }.freeze
-    end
-
-    # Factory method to load from YAML configuration
-    def self.load_from_file(path = nil)
-      path ||= File.join(__dir__, "../../../..", "config", "functions.yaml")
-      functions = Loader.load_file(path)
-      new(functions: functions)
-    end
-
-    def fetch(name, opset: nil)
-      list = @by_name[name] or raise KeyError, "unknown function #{name}"
-      opset ? list.find { |f| f.opset == opset } : list.last
-    end
-
-    # Get function signatures for NEP-20 signature resolution
-    # This bridges RegistryV2 with our existing SignatureResolver
-    def get_function_signatures(name, opset: nil)
-      begin
-        fn = fetch(name, opset: opset)
-        # Convert Signature objects to string representations for NEP-20 parser
-        fn.signatures.map(&:to_signature_string)
-      rescue KeyError
-        []  # Function not found in RegistryV2 - fall back to legacy registry
-      end
-    end
-
-    # Enhanced signature resolution using NEP-20 resolver
-    def choose_signature(fn, arg_shapes)
-      # Use our NEP-20 SignatureResolver for proper dimension handling
-      sig_strings = fn.signatures.map(&:to_signature_string)
-      parsed_sigs = sig_strings.map { |s| SignatureParser.parse(s) }
-      
-      plan = SignatureResolver.choose(signatures: parsed_sigs, arg_shapes: arg_shapes)
-      
-      # Return both the original function signature and the resolution plan
-      {
-        function: fn,
-        signature: plan[:signature],
-        plan: plan
-      }
-    rescue SignatureMatchError => e
-      raise ArgumentError, "no matching signature for #{fn.name} with shapes #{arg_shapes.inspect}: #{e.message}"
-    end
-
-    def resolve_kernel(fn, backend:, conditions: {})
-      ks = fn.kernels.select { |k| k.backend == backend.to_sym }
-      ks = ks.select { |k| conditions.all? { |ck, cv| k.when_&.fetch(ck, cv) == cv } } unless conditions.empty?
-      ks.max_by(&:priority) or raise "no kernel for #{fn.name} backend=#{backend}"
-    end
-
-    # Introspection methods
-    def all_function_names
-      @by_name.keys
-    end
-
-    def function_exists?(name)
-      @by_name.key?(name)
-    end
-
-    def all_functions
-      @by_name.values.flatten
-    end
-  end
-end

data/lib/kumi/core/functions/shape.rb

@@ -1,70 +0,0 @@
-# frozen_string_literal: true
-
-require_relative "dimension"
-
-module Kumi
-  module Core
-    module Functions
-      # Shape utilities with NEP 20 support. A "shape" is an Array<Dimension> of dimension objects.
-      # [] == scalar, [Dimension.new(:i)] == vector along :i, etc.
-      module Shape
-        module_function
-
-        def scalar?(shape) = shape.empty?
-
-        def equal?(a, b) = a.map(&:name) == b.map(&:name)
-
-        # NEP 20 broadcast rules:
-        # - scalar can broadcast to any expected shape
-        # - fixed-size dimensions must match exactly
-        # - broadcastable dimensions with |1 modifier can broadcast against size-1
-        # - flexible dimensions with ? can be omitted if not present in all operands
-        def broadcastable?(got:, expected:)
-          return true if scalar?(got)
-          return false if got.length != expected.length
-
-          got.zip(expected).all? do |got_dim, exp_dim|
-            broadcastable_dimension?(got: got_dim, expected: exp_dim)
-          end
-        end
-
-        def broadcastable_dimension?(got:, expected:)
-          # Same name and modifiers
-          return true if got == expected
-
-          # Same name, different modifiers - check compatibility
-          if got.name == expected.name
-            # Fixed-size dimensions must match exactly
-            if got.fixed_size? || expected.fixed_size?
-              return got.size == expected.size if got.fixed_size? && expected.fixed_size?
-              return false # one fixed, one not - incompatible
-            end
-
-            # Both named dimensions with same name are compatible
-            return true
-          end
-
-          # Different names - only broadcastable with |1 modifier
-          expected.broadcastable? && scalar?([got])
-        end
-
-        # Check if a dimension can be omitted (NEP 20 flexible dimensions)
-        def flexible?(dim)
-          dim.is_a?(Dimension) && dim.flexible?
-        end
-
-        # Check if a dimension can broadcast (NEP 20 broadcastable dimensions)
-        def broadcastable_dimension?(dim)
-          dim.is_a?(Dimension) && dim.broadcastable?
-        end
-
-        # Convenience: find dimensions in set a that are not in set b
-        def dimensions_minus(a, b)
-          a_names = a.map(&:name).to_set
-          b_names = b.map(&:name).to_set
-          (a_names - b_names).to_a
-        end
-      end
-    end
-  end
-end

data/lib/kumi/core/functions/signature.rb

@@ -1,122 +0,0 @@
-# frozen_string_literal: true
-
-module Kumi
-  module Core
-    module Functions
-      # Signature is a small immutable value object describing a function's
-      # vectorization contract with NEP 20 support.
-      #
-      # in_shapes: Array<Array<Dimension>>   e.g., [[Dimension.new(:i)], [Dimension.new(:i)]]
-      # out_shape: Array<Dimension>          e.g., [Dimension.new(:i)], [Dimension.new(:i), Dimension.new(:j)], or []
-      # join_policy: nil | :zip | :product
-      # raw: original string, for diagnostics (optional)
-      class Signature
-        attr_reader :in_shapes, :out_shape, :join_policy, :raw
-
-        def initialize(in_shapes:, out_shape:, join_policy: nil, raw: nil)
-          @in_shapes   = deep_dup(in_shapes).freeze
-          @out_shape   = out_shape.dup.freeze
-          @join_policy = join_policy&.to_sym
-          @raw         = raw
-          validate!
-          freeze
-        end
-
-        def arity = @in_shapes.length
-
-        # Dimensions that appear in any input but not in output (i.e., reduced/dropped).
-        def dropped_axes
-          input_names = @in_shapes.flatten.map(&:name)
-          output_names = @out_shape.map(&:name)
-          (input_names - output_names).uniq.freeze
-        end
-
-        # True if any axis from inputs is dropped (common in aggregates).
-        def reduction?
-          !dropped_axes.empty?
-        end
-
-        def to_h
-          {
-            in_shapes: in_shapes.map(&:dup),
-            out_shape: out_shape.dup,
-            join_policy: join_policy,
-            raw: raw
-          }
-        end
-
-        def inspect
-          "#<Signature #{format_signature}#{" @#{join_policy}" if join_policy}>"
-        end
-
-        def format_signature
-          lhs = in_shapes.map { |dims| "(#{dims.map(&:to_s).join(',')})" }.join(",")
-          rhs = "(#{out_shape.map(&:to_s).join(',')})"
-          "#{lhs}->#{rhs}"
-        end
-
-        # Convert back to string representation for NEP-20 parser compatibility
-        def to_signature_string
-          sig_str = format_signature
-          join_policy ? "#{sig_str}@#{join_policy}" : sig_str
-        end
-
-        private
-
-        def validate!
-          unless @in_shapes.is_a?(Array) && @in_shapes.all? { |s| s.is_a?(Array) }
-            raise SignatureError, "in_shapes must be an array of dimension arrays"
-          end
-
-          @in_shapes.each_with_index do |dims, idx|
-            validate_dimensions!(dims, where: "in_shapes[#{idx}]")
-          end
-
-          validate_dimensions!(@out_shape, where: "out_shape")
-
-          unless [nil, :zip, :product].include?(@join_policy)
-            raise SignatureError, "join_policy must be nil, :zip, or :product; got #{@join_policy.inspect}"
-          end
-
-          # Validate NEP 20 constraints
-          validate_nep20_constraints!
-        end
-
-        def validate_dimensions!(dims, where: "shape")
-          unless dims.is_a?(Array) && dims.all? { |d| d.is_a?(Dimension) }
-            raise SignatureError, "#{where}: must be an array of Dimension objects, got: #{dims.inspect}"
-          end
-
-          # Check for duplicate dimension names within a single argument
-          names = dims.map(&:name)
-          duplicates = names.group_by { |n| n }.select { |_, v| v.size > 1 }.keys
-          raise SignatureError, "#{where}: duplicate dimension names #{duplicates.inspect}" unless duplicates.empty?
-
-          true
-        end
-
-        def validate_nep20_constraints!
-          # Broadcastable dimensions should only appear in inputs, not outputs
-          @out_shape.each do |dim|
-            raise SignatureError, "output dimension #{dim} cannot be broadcastable" if dim.broadcastable?
-          end
-
-          # Fixed-size dimensions in outputs must match corresponding input dimensions
-          all_input_dims = @in_shapes.flatten
-          @out_shape.each do |out_dim|
-            next unless out_dim.fixed_size?
-
-            matching_inputs = all_input_dims.select { |in_dim| in_dim.name == out_dim.name }
-            matching_inputs.each do |in_dim|
-              if in_dim.fixed_size? && in_dim.size != out_dim.size
-                raise SignatureError, "fixed-size dimension #{out_dim.name} has inconsistent sizes: #{in_dim.size} vs #{out_dim.size}"
-              end
-            end
-          end
-        end
-
-        def deep_dup(arr) = arr.map { |x| x.is_a?(Array) ? x.dup : x }
-      end
-    end
-  end
-end

data/lib/kumi/core/functions/signature_parser.rb

@@ -1,86 +0,0 @@
-# frozen_string_literal: true
-
-require_relative "dimension"
-require_relative "signature"
-
-module Kumi
-  module Core
-    module Functions
-      # Parses NEP 20 conformant signature strings like:
-      #  "(),()->()"                    # scalar operations
-      #  "(i),(i)->(i)"                 # vector operations
-      #  "(i),(j)->(i,j)@product"       # matrix operations with join policy
-      #  "(i,j)->(i)"                   # reduction of :j
-      #  "(3),(3)->(3)"                 # fixed-size 3-vectors (cross product)
-      #  "(i?),(i?)->(i?)"               # flexible dimensions
-      #  "(i|1),(i|1)->()"               # broadcastable dimensions
-      #  "(m?,n),(n,p?)->(m?,p?)"        # matmul signature
-      class SignatureParser
-        class << self
-          def parse(str)
-            raise SignatureParseError, "empty signature" if str.nil? || str.strip.empty?
-            lhs, rhs = str.split("->", 2)&.map!(&:strip)
-            raise SignatureParseError, "signature must contain '->': #{str.inspect}" unless rhs
-
-            out_spec, policy = rhs.split("@", 2)&.map!(&:strip)
-            in_shapes = parse_many(lhs)
-            out_shape = parse_axes(out_spec)
-            join_policy = policy&.to_sym
-
-            Signature.new(in_shapes: in_shapes, out_shape: out_shape, join_policy: join_policy, raw: str)
-          rescue SignatureError => e
-            raise
-          rescue StandardError => e
-            raise SignatureParseError, "invalid signature #{str.inspect}: #{e.message}"
-          end
-
-          private
-
-          def parse_many(lhs)
-            # Handle zero arguments case
-            return [] if lhs.strip.empty?
-
-            # split by commas that are *between* groups, not inside them  
-            # simpler approach: split by '),', re-add ')' where needed
-            tokens = lhs.split("),").map { |t| t.strip.end_with?(")") ? t.strip : "#{t.strip})" }
-            tokens.map { |tok| parse_axes(tok) }
-          end
-
-          def parse_axes(spec)
-            spec = spec.strip
-            raise SignatureParseError, "missing parentheses in #{spec.inspect}" unless spec.start_with?("(") && spec.end_with?(")")
-
-            inner = spec[1..-2].strip
-            return [] if inner.empty?
-
-            inner.split(",").map { |dim_str| parse_dimension(dim_str.strip) }
-          end
-
-          # Parse a single dimension with NEP 20 modifiers
-          # Examples: "i", "3", "n?", "i|1"
-          def parse_dimension(dim_str)
-            return Dimension.new(:empty) if dim_str.empty?
-
-            # Extract modifiers
-            flexible = dim_str.end_with?("?")
-            dim_str = dim_str.chomp("?") if flexible
-
-            broadcastable = dim_str.end_with?("|1")
-            dim_str = dim_str.chomp("|1") if broadcastable
-
-            # Parse name (symbol or integer)
-            name = if dim_str.match?(/^\d+$/)
-                     dim_str.to_i
-                   else
-                     dim_str.to_sym
-                   end
-
-            Dimension.new(name, flexible: flexible, broadcastable: broadcastable)
-          rescue StandardError => e
-            raise SignatureParseError, "invalid dimension #{dim_str.inspect}: #{e.message}"
-          end
-        end
-      end
-    end
-  end
-end

data/lib/kumi/core/functions/signature_resolver.rb

@@ -1,272 +0,0 @@
-# frozen_string_literal: true
-
-require_relative "errors"
-require_relative "shape"
-require_relative "signature"
-
-module Kumi
-  module Core
-    module Functions
-      # Given a set of signatures and actual argument shapes, pick the best match.
-      # Supports NEP 20 extensions: fixed-size, flexible, and broadcastable dimensions.
-      #
-      # Inputs:
-      #   signatures : Array<Signature> (with Dimension objects)
-      #   arg_shapes : Array<Array<Symbol|Integer>>   e.g., [[:i], [:i]] or [[], [3]] or [[2, :i]]
-      #
-      # Returns:
-      #   { signature:, result_axes:, join_policy:, dropped_axes:, effective_signature: }
-      #
-      # NEP 20 Matching rules:
-      # - Arity must match exactly (before flexible dimension resolution).
-      # - Fixed-size dimensions (integers) must match exactly.
-      # - Flexible dimensions (?) can be omitted if not present in all operands.
-      # - Broadcastable dimensions (|1) can match scalar or size-1 dimensions.
-      # - For each param position, shapes are checked according to NEP 20 rules.
-      # - We prefer exact matches, then flexible matches, then broadcast matches.
-      class SignatureResolver
-        class << self
-          def choose(signatures:, arg_shapes:)
-            # Handle empty arg_shapes for zero-arity functions
-            arg_shapes = [] if arg_shapes.nil?
-            sanity_check_args!(arg_shapes)
-
-            candidates = signatures.map do |sig|
-              score = match_score(sig, arg_shapes)
-              next if score.nil?
-
-              # Convert arg_shapes to normalized Dimension arrays for environment building
-              normalized_args = arg_shapes.map { |shape| normalize_shape(shape) }
-              env = build_dimension_environment(sig, normalized_args)
-              next if env.nil?  # Skip candidates with dimension conflicts
-
-              {
-                signature: sig,
-                score: score,
-                result_axes: sig.out_shape.map(&:name), # Convert Dimension objects to names for backward compatibility
-                join_policy: sig.join_policy,
-                dropped_axes: sig.dropped_axes.map { |name| name.is_a?(Symbol) ? name : name.to_sym }, # Convert to symbols
-                env: env
-              }
-            end.compact
-
-            raise SignatureMatchError, mismatch_message(signatures, arg_shapes) if candidates.empty?
-
-            # Lower score is better: 0 = exact-everywhere, then number of broadcasts
-            best = candidates.min_by { |c| c[:score] }
-            
-            # Add effective signature and environment for analyzer/lowering
-            best[:effective_signature] = {
-              in_shapes: best[:signature].in_shapes.map { |dims| dims.map(&:name) },
-              out_shape: best[:signature].out_shape.map(&:name),
-              join_policy: best[:signature].join_policy
-            }
-            # env is already included from candidate building
-            
-            best
-          end
-
-          private
-
-          def sanity_check_args!(arg_shapes)
-            unless arg_shapes.is_a?(Array) &&
-                   arg_shapes.all? { |s| s.is_a?(Array) && s.all? { |a| a.is_a?(Symbol) || a.is_a?(Integer) } }
-              raise SignatureMatchError, "arg_shapes must be an array of dimension arrays (symbols or integers), got: #{arg_shapes.inspect}"
-            end
-          end
-
-          # Returns an integer "broadcast cost" or nil if not matchable.
-          # Lower score = better match: 0 = exact, then increasing cost for broadcasts/flexibility
-          def match_score(sig, arg_shapes)
-            return nil unless sig.arity == arg_shapes.length
-
-            # Convert arg_shapes to normalized Dimension arrays for comparison
-            normalized_args = arg_shapes.map { |shape| normalize_shape(shape) }
-
-            # Try to match each argument against its expected signature shape
-            cost = 0
-            sig.in_shapes.each_with_index do |expected_dims, idx|
-              got_dims = normalized_args[idx]
-              arg_cost = match_argument_cost(got: got_dims, expected: expected_dims)
-              return nil if arg_cost.nil?
-
-              cost += arg_cost
-            end
-
-            # Additional checks for join_policy constraints
-            return nil unless valid_join_policy?(sig, normalized_args)
-
-            cost
-          end
-
-          private
-
-          # Convert a shape array (symbols/integers) to normalized Dimension array
-          def normalize_shape(shape)
-            shape.map do |dim|
-              case dim
-              when Symbol
-                Dimension.new(dim)
-              when Integer
-                Dimension.new(dim)
-              else
-                raise SignatureMatchError, "Invalid dimension type: #{dim.class}"
-              end
-            end
-          end
-
-          # Calculate cost of matching one argument against expected dimensions
-          def match_argument_cost(got:, expected:)
-            # Handle scalar first
-            if got.empty?
-              return expected.empty? ? 0 : (expected.any?(&:flexible?) ? 10 : 1) # scalar broadcast or flexible-tail
-            end
-
-            # Try strict matching first if no flexible dimensions
-            if !expected.any?(&:flexible?) && got.length == expected.length
-              total = 0
-              got.zip(expected).each do |g, e|
-                c = match_dimension_cost(got: g, expected: e)
-                return nil if c.nil?
-                total += c
-              end
-              return total
-            end
-
-            # Use right-aligned flexible matching
-            right_align_match(got: got, expected: expected)
-          end
-
-          # Right-aligned matching for flexible dimensions (NEP 20 ? modifier)
-          def right_align_match(got:, expected:)
-            gi = got.length - 1
-            ei = expected.length - 1
-            cost = 0
-
-            while ei >= 0
-              exp = expected[ei]
-
-              if exp.flexible? && gi < 0
-                # optional tail dimension that we don't have → ok, consume expected only
-                ei -= 1
-                cost += 10
-                next
-              end
-
-              return nil if gi < 0 # ran out of got dims and exp wasn't flexible
-
-              got_dim = got[gi]
-              dim_cost = match_dimension_cost(got: got_dim, expected: exp)
-              if dim_cost.nil?
-                # if exp is flexible, we can try to drop it
-                if exp.flexible?
-                  ei -= 1
-                  cost += 10
-                  next
-                else
-                  return nil
-                end
-              else
-                cost += dim_cost
-                gi -= 1
-                ei -= 1
-              end
-            end
-
-            # if we still have leftover got dims, argument is longer than expected → not a match
-            return nil if gi >= 0
-
-            cost
-          end
-
-          # Calculate cost of matching one dimension against another
-          def match_dimension_cost(got:, expected:)
-            return 0 if got == expected # Exact match
-
-            # Fixed-size equality
-            if got.fixed_size? && expected.fixed_size?
-              return got.size == expected.size ? 0 : nil
-            end
-
-            # Same symbolic name (ignoring modifiers) → ok unless one is fixed and the other isn't (penalize)
-            if got.named? && expected.named? && got.name == expected.name
-              return (got.fixed_size? || expected.fixed_size?) ? 2 : 0
-            end
-
-            # Broadcastable expected dim accepts scalar or size-1
-            if expected.broadcastable?
-              # scalar at argument level would have been handled in match_argument_cost
-              # so here we check for size-1 fixed dimensions
-              return 3 if got.fixed_size? && got.size == 1
-              # Named dimensions that could be size-1 at runtime also get broadcast cost
-              return 3 if got.named?
-            end
-
-            nil # No match possible
-          end
-
-          # Check if join_policy constraints are satisfied
-          def valid_join_policy?(sig, normalized_args)
-            return true if sig.join_policy # :zip or :product allows different axes
-
-            # nil join_policy: check if dimension names are consistent
-            non_scalar_args = normalized_args.reject { |a| Shape.scalar?(a) }
-            return true if non_scalar_args.empty?
-
-            # For nil join_policy, we allow different dimension names if:
-            # 1. All args have same dimension names (element-wise operations), OR
-            # 2. The constraint solver can validate cross-dimensional consistency (like matmul)
-            first_names = non_scalar_args.first.map(&:name)
-            same_names = non_scalar_args.all? { |arg| arg.map(&:name) == first_names }
-            
-            return true if same_names
-            
-            # If dimension names differ, check if constraint solver can handle it
-            # This allows operations like matmul where dimensions are linked across arguments
-            env = build_dimension_environment(sig, normalized_args)
-            !env.nil?
-          end
-
-          def mismatch_message(signatures, arg_shapes)
-            sigs = signatures.map(&:inspect).join(", ")
-            "no matching signature for shapes #{pp_shapes(arg_shapes)} among [#{sigs}]"
-          end
-
-          def pp_shapes(shapes)
-            shapes.map { |ax| "(#{ax.join(',')})" }.join(", ")
-          end
-
-          # Build dimension environment by checking consistency of named dimensions across arguments
-          def build_dimension_environment(sig, normalized_args)
-            env = {}
-            
-            # Walk all expected dimensions across all arguments
-            sig.in_shapes.each_with_index do |expected_shape, arg_idx|
-              got_shape = normalized_args[arg_idx] || []
-              
-              expected_shape.each_with_index do |exp_dim, dim_idx|
-                next unless exp_dim.named? && dim_idx < got_shape.length
-                
-                got_dim = got_shape[dim_idx]
-                dim_name = exp_dim.name
-                
-                # Check for consistency: same dimension name must map to same concrete value
-                if env.key?(dim_name)
-                  # If we've seen this dimension name before, it must match
-                  if env[dim_name] != got_dim
-                    return nil  # Inconsistent binding - signature doesn't match
-                  end
-                else
-                  # First time seeing this dimension name - record the binding
-                  env[dim_name] = got_dim
-                end
-              end
-            end
-            
-            env
-          end
-        end
-      end
-    end
-  end
-end