Sha256: eac2aa6d7061318f9bbd884f298f645fe1c58be83dc8b9b066194b198ac75527
Contents?: true
Size: 577 Bytes
Versions: 9
Compression:
Stored size: 577 Bytes
Contents
TensorStream::OpMaker.define_operation :argmax do |op|
op.what_it_does "Returns the index with the largest value across axes of a tensor."
op.parameter :input_a, "tensor X", validate: 'NUMERIC_TYPES'
op.parameter :axis, "Describes which axis of the input tensor to reduce across. For vectors, use axis = 0", :nil, validate: 'INTEGER_TYPES'
op.option :name, "Optional name", :nil
op.option :dimension, "Same as axis", :nil
op.option :output_type, "Output data type defaults to int32", ":int32"
op.define_gradient do |grad, node, params|
[nil, nil]
end
end
Version data entries
9 entries across 9 versions & 1 rubygems