You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.

stack.py 1.7 kB

1234567891011121314151617181920212223242526272829303132333435363738
  1. # Copyright 2021 Huawei Technologies Co., Ltd
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. # ============================================================================
  15. """Stack op"""
  16. from mindspore.ops.op_info_register import op_info_register, CpuRegOp, DataType
  17. stack_op_info = CpuRegOp("Stack") \
  18. .input(0, "x", "required") \
  19. .output(0, "y", "required") \
  20. .dtype_format(DataType.I8_Default, DataType.I8_Default) \
  21. .dtype_format(DataType.I16_Default, DataType.I16_Default) \
  22. .dtype_format(DataType.I32_Default, DataType.I32_Default) \
  23. .dtype_format(DataType.I64_Default, DataType.I64_Default) \
  24. .dtype_format(DataType.U8_Default, DataType.U8_Default) \
  25. .dtype_format(DataType.U16_Default, DataType.U16_Default) \
  26. .dtype_format(DataType.U32_Default, DataType.U32_Default) \
  27. .dtype_format(DataType.U64_Default, DataType.U64_Default) \
  28. .dtype_format(DataType.F16_Default, DataType.F16_Default) \
  29. .dtype_format(DataType.F32_Default, DataType.F32_Default) \
  30. .dtype_format(DataType.BOOL_Default, DataType.BOOL_Default) \
  31. .get_op_info()
  32. @op_info_register(stack_op_info)
  33. def _stack_cpu():
  34. """Stack cpu register"""
  35. return