diff --git a/build.sh b/build.sh
index 3fb4cbebd193061a15fbcca2f4b3a5255c1ed318..0fbe4228b461ef1f227d207d97d227765d6c7250 100644
--- a/build.sh
+++ b/build.sh
@@ -1,30 +1,39 @@
 #!/bin/bash
 
-# 配置tensorflow .bazlerc选项
-printf "\n\nn\nn\nn\nn\nn\nn\n\nn\n" | ./configure
-
-# install ktfop
-mkdir -p third_party/ktfop/include
-mkdir -p third_party/ktfop/lib
-cp -a /usr/local/sra_inference/lib/neon/libktfop* third_party/ktfop/lib/
-cp -a /usr/local/sra_inference/include/* third_party/ktfop/include/
-
 # 使用 rpm -qi 提取KML Version 信息
 VERSION=$(rpm -qi kml | grep "^Version" | awk '{print $3}')
 # 输出 Version
 echo "KML Version: $VERSION"
+KTFOP_LIB_PATH=/usr/local/sra_inference/lib/neon
+KTFOP_INCLUDE_PATH=/usr/local/sra_inference/include
+KBLAS_LIB_PATH=""
+KBLAS_INCLUDE_PATH=/usr/local/kml/include
 
-# 比较版本号是否 >= 2.5.0
-if [[ $(printf '%s\n' "$VERSION" "2.5.0" | sort -V | head -n1) == "2.5.0" ]]; then 
-    LD_LIBRARY_PATH=/usr/local/sra_inference/lib/neon:/usr/local/kml/lib/neon/kblas/locking/:$LD_LIBRARY_PATH \
-    CPLUS_INCLUDE_PATH=/usr/local/kml/include:$CPLUS_INCLUDE_PATH \
-    C_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:$C_INCLUDE_PATH \
-    bazel build --config=opt --cxxopt="-D_GLIBCXX_USE_CXX11_ABI=0" --verbose_failures //tensorflow/tools/pip_package:build_pip_package && \
-    ./bazel-bin/tensorflow/tools/pip_package/build_pip_package ./out
+if [[ $(printf '%s\n' "$VERSION" "2.5.0" | sort -V | head -n1) == "2.5.0" ]]; then
+    KBLAS_LIB_PATH=/usr/local/kml/lib/neon/kblas/locking
 else
-    LD_LIBRARY_PATH=/usr/local/sra_inference/lib/neon:/usr/local/kml/lib/kblas/locking/:$LD_LIBRARY_PATH \
-    CPLUS_INCLUDE_PATH=/usr/local/kml/include:$CPLUS_INCLUDE_PATH \
-    C_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:$C_INCLUDE_PATH \
-    bazel build --config=opt --cxxopt="-D_GLIBCXX_USE_CXX11_ABI=0" --verbose_failures //tensorflow/tools/pip_package:build_pip_package && \
-    ./bazel-bin/tensorflow/tools/pip_package/build_pip_package ./out
+    KBLAS_LIB_PATH=/usr/local/kml/lib/kblas/locking
+fi
+
+if [ -n "$KBLAS_LIB_DIR" ]; then
+    KBLAS_LIB_PATH="$KBLAS_LIB_DIR"
 fi
+
+echo "ktfop lib path: " $KTFOP_LIB_PATH
+echo "kblas lib path: " $KBLAS_LIB_PATH
+
+# install ktfop and kblas
+mkdir -p third_party/ktfop/include
+mkdir -p third_party/ktfop/lib
+cp -a ${KTFOP_LIB_PATH}/libktfop* third_party/ktfop/lib/
+cp -a ${KTFOP_INCLUDE_PATH}/ktfop.h third_party/ktfop/include/
+cp -a ${KBLAS_LIB_PATH}/* third_party/ktfop/lib/
+cp -a ${KBLAS_INCLUDE_PATH}/kblas.h third_party/ktfop/include/
+
+# 配置tensorflow .bazlerc选项
+printf "\n\nn\nn\nn\nn\nn\nn\n\nn\n" | ./configure
+LD_LIBRARY_PATH=${KTFOP_LIB_PATH}:${KBLAS_LIB_PATH}:$LD_LIBRARY_PATH \
+CPLUS_INCLUDE_PATH=${KBLAS_INCLUDE_PATH}:$CPLUS_INCLUDE_PATH \
+C_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:$C_INCLUDE_PATH \
+bazel build --config=opt --cxxopt="-D_GLIBCXX_USE_CXX11_ABI=0" --verbose_failures //tensorflow/tools/pip_package:build_pip_package && \
+./bazel-bin/tensorflow/tools/pip_package/build_pip_package ./out
diff --git a/tensorflow/python/lib/core/bfloat16.cc b/tensorflow/python/lib/core/bfloat16.cc
index 7d170113a3e782e8c4e8fc952e01a3095c8c6e80..fde3a83770280038b777a141693d117dace4b41f 100644
--- a/tensorflow/python/lib/core/bfloat16.cc
+++ b/tensorflow/python/lib/core/bfloat16.cc
@@ -490,7 +490,7 @@ bool RegisterBfloat16Cast(int numpy_type, bool cast_is_safe) {
 }
 
 template <typename InType, typename OutType, typename Functor>
-void BinaryUFunc(char** args, npy_intp const* dimensions, npy_intp const* steps,
+void BinaryUFunc(char** args, npy_intp* dimensions, npy_intp* steps,
                  void* data) {
   const char* i0 = args[0];
   const char* i1 = args[1];
@@ -506,7 +506,7 @@ void BinaryUFunc(char** args, npy_intp const* dimensions, npy_intp const* steps,
 }
 
 template <typename Functor>
-void CompareUFunc(char** args, npy_intp const* dimensions, npy_intp const* steps,
+void CompareUFunc(char** args, npy_intp* dimensions, npy_intp* steps,
                   void* data) {
   BinaryUFunc<bfloat16, npy_bool, Functor>(args, dimensions, steps, data);
 }
diff --git a/third_party/ktfop/BUILD b/third_party/ktfop/BUILD
index bfa1f6e1d76e66b73d48e1e92f88bf289ee93d4b..a9759d03467b82c7ae57b9a3c01031bb8811cd56 100644
--- a/third_party/ktfop/BUILD
+++ b/third_party/ktfop/BUILD
@@ -1,4 +1,15 @@
-licenses(["notice"]) 
+licenses(["notice"])
+
+cc_library(
+    name = "libkblas",
+    srcs = ["lib/libkblas.so"],
+    hdrs = [
+        "include/kblas.h",
+    ],
+    includes = ["."],
+    visibility = ["//visibility:public"],
+)
+
 cc_library(
     name = "libktfop",
     srcs = ["lib/libktfop.so"],
@@ -8,5 +19,6 @@ cc_library(
         "math/tools.h",
     ],
     includes = ["."],
+    deps = [":libkblas"],
     visibility = ["//visibility:public"],
 )