com.yahoo.tensor.functions.Range 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.TensorType;
import com.yahoo.tensor.evaluation.Name;

import java.util.List;
import java.util.Objects;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * A tensor generator which returns a tensor of any dimension filled with the sum of the tensor
 * indexes of each position.
 *
 * @author bratseth
 */
public class Range extends CompositeTensorFunction {

    private final TensorType type;
    private final Function, Double> rangeFunction;

    public Range(TensorType type) {
        this.type = type;
        this.rangeFunction = ScalarFunctions.sum(dimensionNames().toList());
    }

    @Override
    public List> arguments() { return List.of(); }

    @Override
    public TensorFunction withArguments(List> arguments) {
        if ( arguments.size() != 0)
            throw new IllegalArgumentException("Range must have 0 arguments, got " + arguments.size());
        return this;
    }

    @Override
    public PrimitiveTensorFunction toPrimitive() {
        return new Generate<>(type, rangeFunction);
    }

    @Override
    public String toString(ToStringContext context) {
        return "range(" + type + ")";
    }

    private Stream dimensionNames() {
        return type.dimensions().stream().map(TensorType.Dimension::name);
    }

    @Override
    public int hashCode() {
        return Objects.hash("range", type, rangeFunction);
    }

}