com.yahoo.tensor.functions.DenseSubspaceFunction Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of vespajlib Show documentation

Library for use in Java components of Vespa. Shared code which do not fit anywhere else.

There is a newer version: 8.441.21

// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.tensor.functions;

import com.yahoo.tensor.Tensor;
import com.yahoo.tensor.TensorType;
import com.yahoo.tensor.evaluation.EvaluationContext;
import com.yahoo.tensor.evaluation.MapEvaluationContext;
import com.yahoo.tensor.evaluation.Name;
import com.yahoo.tensor.evaluation.TypeContext;

import java.util.Optional;
import java.util.function.Function;

/**
 * A function suitable for use in MapSubspaces
 *
 * @author arnej
 */
class DenseSubspaceFunction {

    private final String argName;
    private final TensorFunction function;

    public DenseSubspaceFunction(String argName, TensorFunction function) {
        this.argName = argName;
        this.function = function;
    }

    Tensor map(Tensor subspace) {
        var context = new MapEvaluationContext();
        context.put(argName, subspace);
        return function.evaluate(context);
    }

    class MyTypeContext implements TypeContext {
        private final TensorType subspaceType;
        MyTypeContext(TensorType subspaceType) { this.subspaceType = subspaceType; }
        public TensorType getType(NAMETYPE name) { return getType(name.name()); }
        public TensorType getType(String name) { return argName.equals(name) ? subspaceType : null; }
        public String resolveBinding(String name) { return name; }
    }

    TensorType outputType(TensorType subspaceType) {
        var context = new MyTypeContext(subspaceType);
        var result = function.type(context);
        if (result.mappedSubtype().rank() > 0) {
            throw new IllegalArgumentException("function used in map_subspaces type had mapped dimensions: " + result);
        }
        return result;
    }

    public String toString() {
        return "f(" + argName + ")(" + function + ")";
    }

}