blob: b6655a153616d1c4b748a6a3a7ed59df13c1f1d4 (
plain) (
blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
|
// 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<NAMETYPE extends Name> {
private final String argName;
private final TensorFunction<NAMETYPE> function;
public DenseSubspaceFunction(String argName, TensorFunction<NAMETYPE> function) {
this.argName = argName;
this.function = function;
}
Tensor map(Tensor subspace) {
var context = new MapEvaluationContext<NAMETYPE>();
context.put(argName, subspace);
return function.evaluate(context);
}
class MyTypeContext implements TypeContext<NAMETYPE> {
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; }
}
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 + ")";
}
}
|