org.nd4j.linalg.api.ops.impl.shape.Transpose Maven / Gradle / Ivy
The newest version!
/*
* ******************************************************************************
* *
* *
* * This program and the accompanying materials are made available under the
* * terms of the Apache License, Version 2.0 which is available at
* * https://www.apache.org/licenses/LICENSE-2.0.
* *
* * See the NOTICE file distributed with this work for additional
* * information regarding copyright ownership.
* * Unless required by applicable law or agreed to in writing, software
* * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* * License for the specific language governing permissions and limitations
* * under the License.
* *
* * SPDX-License-Identifier: Apache-2.0
* *****************************************************************************
*/
package org.nd4j.linalg.api.ops.impl.shape;
import org.nd4j.shade.guava.primitives.Ints;
import lombok.val;
import onnx.Onnx;
import org.nd4j.autodiff.samediff.SDVariable;
import org.nd4j.autodiff.samediff.SameDiff;
import org.nd4j.autodiff.samediff.VariableType;
import org.nd4j.common.base.Preconditions;
import org.nd4j.imports.descriptors.properties.PropertyMapping;
import org.nd4j.imports.graphmapper.tf.TFGraphMapper;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.api.ops.DynamicCustomOp;
import org.nd4j.linalg.exception.ND4JIllegalStateException;
import org.nd4j.common.util.ArrayUtil;
import org.tensorflow.framework.AttrValue;
import org.tensorflow.framework.GraphDef;
import org.tensorflow.framework.NodeDef;
import java.util.*;
public class Transpose extends DynamicCustomOp {
protected int[] permuteDims;
public Transpose(SameDiff sameDiff, SDVariable i_v) {
super(null, sameDiff, new SDVariable[]{i_v});
}
public Transpose(SameDiff sameDiff, SDVariable in, int[] permuteDims){
super(null, sameDiff, new SDVariable[]{in});
this.permuteDims = permuteDims;
}
protected Transpose(SameDiff sameDiff, SDVariable in, SDVariable permuteDims){
super(null, sameDiff, new SDVariable[]{in, permuteDims});
}
public Transpose(INDArray input, INDArray result){
super(null, new INDArray[]{input}, result == null ? null : new INDArray[]{result}, null, (List) null);
}
public Transpose(INDArray input){
this(input, null);
}
public Transpose() {
}
@Override
public Map> mappingsForFunction() {
Map> ret = new LinkedHashMap<>();
Map map = new LinkedHashMap<>();
val mapping = PropertyMapping.builder()
.onnxAttrName("perm")
.propertyNames(new String[]{"permuteDims"})
.tfInputPosition(1)
.build();
map.put("permuteDims", mapping);
ret.put(tensorflowName(), map);
ret.put(onnxName(), map);
return ret;
}
@Override
public String opName() {
return "transpose";
}
@Override
public String onnxName() {
return "Transpose";
}
@Override
public String tensorflowName() {
return "Transpose";
}
@Override
public void initFromTensorFlow(NodeDef nodeDef, SameDiff initWith, Map attributesForNode, GraphDef graph) {
super.initFromTensorFlow(nodeDef, initWith, attributesForNode, graph);
//permute dimensions are not specified as second input
if (nodeDef.getInputCount() < 2)
return;
NodeDef permuteDimsNode = null;
for (int i = 0; i < graph.getNodeCount(); i++) {
if (graph.getNode(i).getName().equals(nodeDef.getInput(1))) {
permuteDimsNode = graph.getNode(i);
}
}
INDArray permuteArrayOp = TFGraphMapper.getNDArrayFromTensor(permuteDimsNode);
if (permuteArrayOp != null) {
this.permuteDims = permuteArrayOp.data().asInt();
}
//handle once properly mapped
if (arg().getShape() == null || arg().getVariableType() == VariableType.PLACEHOLDER || arg().getArr() == null) {
return;
}
INDArray arr = sameDiff.getArrForVarName(arg().name());
if(permuteArrayOp != null){
addInputArgument(arr, permuteArrayOp);
} else {
addInputArgument(arr);
}
if (arr != null && permuteDims == null) {
this.permuteDims = ArrayUtil.reverseCopy(ArrayUtil.range(0, arr.rank()));
}
if (permuteDims != null && permuteDims.length < arg().getShape().length)
throw new ND4JIllegalStateException("Illegal permute found. Not all dimensions specified");
}
@Override
public void initFromOnnx(Onnx.NodeProto node, SameDiff initWith, Map attributesForNode, Onnx.GraphProto graph) {
if (!attributesForNode.containsKey("perm")) {
} else
this.permuteDims = Ints.toArray(attributesForNode.get("perm").getIntsList());
}
@Override
public List doDiff(List i_v) {
SDVariable ret;
if(permuteDims == null) {
ret = sameDiff.transpose(i_v.get(0));
} else {
int[] reverse = ArrayUtil.invertPermutation(permuteDims);
ret = sameDiff.permute(i_v.get(0), reverse);
}
return Collections.singletonList(ret);
}
@Override
public List calculateOutputDataTypes(List dataTypes){
Preconditions.checkState(dataTypes != null && (dataTypes.size() == 1 || dataTypes.size() == 2),
"Expected list with 1 or 2 datatype for %s, got %s", getClass(), dataTypes);
if(dArguments != null && !dArguments.isEmpty())
return Collections.singletonList(dArguments.get(0));
//Output type is same as input type. Second input is permute dimensions as array
return Collections.singletonList(dataTypes.get(0));
}
}